import re
import gradio as gr
import spaces
import torch
from threading import Thread
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer

MODEL_ID = "CMSManhattan/JiRackDeltaNet_27b"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, dtype=torch.bfloat16, trust_remote_code=True
)
model.to("cuda")

# Останавливаемся на конце реплики ассистента
STOP_IDS = {tokenizer.eos_token_id}
for tok in ("<|im_end|>", "<|endoftext|>"):
    tid = tokenizer.convert_tokens_to_ids(tok)
    if tid is not None and tid != tokenizer.unk_token_id:
        STOP_IDS.add(tid)
STOP_IDS = [i for i in STOP_IDS if i is not None]

THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL)


def clean(text):
    text = THINK_RE.sub("", text)
    text = text.replace("<think>", "").replace("</think>", "")
    # если модель всё же начала выдуманную новую реплику, отрезаем её
    text = re.split(r"\n(?:user|assistant)\n", text)[0]
    return text.strip()


def as_text(content):
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        return "".join(
            p.get("text", "") if isinstance(p, dict) else str(p) for p in content
        )
    return str(content)


def build_messages(message, history):
    msgs = []
    for turn in history:
        if isinstance(turn, dict):
            msgs.append({"role": turn["role"], "content": as_text(turn["content"])})
        else:
            user, bot = turn
            msgs.append({"role": "user", "content": as_text(user)})
            if bot:
                msgs.append({"role": "assistant", "content": as_text(bot)})
    msgs.append({"role": "user", "content": as_text(message)})
    return msgs


@spaces.GPU(size="xlarge", duration=120)
def respond(message, history, max_new_tokens=512, temperature=0.7):
    inputs = tokenizer.apply_chat_template(
        build_messages(message, history),
        add_generation_prompt=True,
        return_tensors="pt",
        return_dict=True,
    ).to("cuda")
    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )
    Thread(
        target=model.generate,
        kwargs=dict(
            **inputs,
            streamer=streamer,
            max_new_tokens=int(max_new_tokens),
            do_sample=True,
            temperature=float(temperature),
            top_p=0.8,
            top_k=20,
            eos_token_id=STOP_IDS,
        ),
    ).start()
    out = ""
    for piece in streamer:
        out += piece
        yield clean(out)


demo = gr.ChatInterface(
    respond,
    title="JiRack DeltaNet 27B",
    description="Try the JiRack DeltaNet 27B model.",
    additional_inputs=[
        gr.Slider(64, 1024, value=512, step=64, label="Max new tokens"),
        gr.Slider(0.1, 1.5, value=0.7, step=0.1, label="Temperature"),
    ],
)
demo.launch()