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

#Qwen/Qwen2.5-14B-Instruct-1M
#Qwen/Qwen2-0.5B
model_name = "bobber/Qwen-0.5B-GRPO"
subfolder = "Qwen-0.5B-GRPO/checkpoint-1868"

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    subfolder=subfolder,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name, subfolder=subfolder)
SYSTEM_PROMPT = """
Respond in the following format:
<reasoning>
...
</reasoning>
<answer>
...
</answer>
"""

@spaces.GPU
def generate(prompt, history):
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": prompt}
    ]
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
    model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
    
    generated_ids = model.generate(
        **model_inputs,
        max_new_tokens=512
    )
    generated_ids = [
        output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
    ]
    
    response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
    return response



chat_interface = gr.ChatInterface(
    fn=generate,
)
chat_interface.launch(share=True)