import spaces
import torch
import gradio as gr
from threading import Thread
from transformers import AutoTokenizer, Qwen3_5ForConditionalGeneration, TextIteratorStreamer
from torchao.quantization import quantize_, Int8WeightOnlyConfig

MODEL_ID = "UnstableLlama/Semancer-27B"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = Qwen3_5ForConditionalGeneration.from_pretrained(
    MODEL_ID,
    dtype=torch.bfloat16,
).to("cuda")
quantize_(model, Int8WeightOnlyConfig())  # ~27 GB, fits the large slice
model.eval()


def _estimate_duration(message, history, system_prompt="", max_new_tokens=1024, *args, **kwargs):
    # ~10 tok/s worst case for int8-weight-only 27B
    return min(240, 30 + int(max_new_tokens) // 10)


@spaces.GPU(duration=_estimate_duration)
def chat(
    message: str,
    history: list,
    system_prompt: str = "",
    max_new_tokens: int = 1024,
    temperature: float = 0.95,
    top_p: float = 1.0,
    top_k: int = 50,
    min_p: float = 0.04,
):
    """Chat with Semancer-27B, an occult philosophy fine-tune of Qwen3.6 27B.

    Ask it about truth, energy, entropy, free will, justice — it answers
    from first principles in prose that sounds nothing like a standard LLM.

    Args:
        message: The user's message.
        history: Chat history as a list of {"role", "content"} dicts.
        system_prompt: Optional system prompt to guide the model.
        max_new_tokens: Maximum number of tokens to generate.
        temperature: Sampling temperature.
        top_p: Nucleus sampling probability threshold.
        top_k: Top-k sampling limit.
        min_p: Minimum token probability, scaled by the top token's probability.
    """
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    for m in history:
        messages.append(m)
    messages.append({"role": "user", "content": message})

    inputs = tokenizer.apply_chat_template(
        messages,
        add_generation_prompt=True,
        return_tensors="pt",
        tokenize=True,
        enable_thinking=False,
    )
    if hasattr(inputs, "input_ids"):
        input_ids = inputs.input_ids
    else:
        input_ids = inputs
    input_ids = input_ids.to("cuda")

    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )

    generation_kwargs = dict(
        input_ids=input_ids,
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        temperature=temperature,
        top_p=top_p,
        top_k=top_k,
        min_p=min_p,
        pad_token_id=tokenizer.pad_token_id,
    )

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    partial = ""
    for chunk in streamer:
        partial += chunk
        yield partial

    thread.join()


CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

demo = gr.ChatInterface(
    fn=chat,
    title="🔮 Semancer-27B",
    description=(
        "An occult philosophy fine-tune of "
        "[Qwen3.6-27B](https://huggingface.co/Qwen/Qwen3.6-27B) "
        "by [UnstableLlama](https://huggingface.co/UnstableLlama).\n\n"
        "Ask it about truth, energy, entropy, free will, justice, intelligence "
        "— it answers from first principles you have not encountered, "
        "in prose that sounds nothing like your inference stack."
    ),
    examples=[
        ["Are you conscious?", "", 1024, 0.95, 1.0, 50, 0.04],
        ["What is truth?", "", 1024, 0.95, 1.0, 50, 0.04],
        ["What is forgiveness?", "", 1024, 0.95, 1.0, 50, 0.04],
        ["Explain entropy from first principles.", "", 1024, 0.95, 1.0, 50, 0.04],
        ["Does the Chinese room understand anything?", "", 1024, 0.95, 1.0, 50, 0.04],
    ],
    additional_inputs=[
        gr.Textbox(
            value="",
            label="System prompt",
            lines=2,
            placeholder="Optional — leave blank for no system prompt",
        ),
        gr.Slider(64, 4096, value=1024, step=64, label="Max new tokens"),
        gr.Slider(0.01, 2.0, value=0.95, step=0.01, label="Temperature"),
        gr.Slider(0.1, 1.0, value=1.0, step=0.05, label="Top-p"),
        gr.Slider(1, 256, value=50, step=1, label="Top-k"),
        gr.Slider(0.0, 0.5, value=0.04, step=0.01, label="Min-p"),
    ],
    additional_inputs_accordion=gr.Accordion("⚙️ Advanced settings", open=False),
    cache_examples=True,
    cache_mode="lazy",
)

demo.launch(mcp_server=True, theme=gr.themes.Citrus(), css=CSS)