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 = "bartowski/simplescaling_s1-32B-GGUF"
subfolder = "Qwen-0.5B-GRPO/checkpoint-1868"
filename = "simplescaling_s1-32B-Q6_K_L.gguf"
torch_dtype = torch.float32 # could be torch.float16 or torch.bfloat16 too

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    # subfolder=subfolder,
    gguf_file=filename,
    torch_dtype=torch_dtype,
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name
    , gguf_file=filename
    # , 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)