import gradio as gr
import torch
import spaces
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO = "prathamkode/particle-1.0"
tok = AutoTokenizer.from_pretrained(REPO)
model = AutoModelForCausalLM.from_pretrained(
    REPO,
    torch_dtype=torch.float32,
    device_map="cpu",
)
model.eval()

@spaces.GPU
def _zero_gpu_placeholder():
    return None

def chat(message, history):
    messages = [{"role": m["role"], "content": m["content"]} for m in history]
    messages.append({"role": "user", "content": message})
    prompt = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    ids = tok(prompt, return_tensors="pt")
    ids.pop("token_type_ids", None)
    out = model.generate(
        **ids,
        max_new_tokens=64,
        do_sample=True,
        temperature=0.7,
        top_k=50,
        pad_token_id=tok.pad_token_id,
        eos_token_id=tok.eos_token_id,
    )
    text = tok.decode(out[0], skip_special_tokens=False)
    if "<|assistant|>" in text:
        text = text.split("<|assistant|>")[-1]
    return text.replace("<|endoftext|>", "").replace("<|padding|>", "").strip()


demo = gr.ChatInterface(chat, title="particle-1.0")

if __name__ == "__main__":
    demo.launch(ssr_mode=False)