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

model_name = "brucoder/WINTER-FROST-2-PRO"

print("Loading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(model_name)

print("Loading model (this may take a while on CPU)...")
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    low_cpu_mem_usage=True,
)
model.eval()

def predict(prompt: str, max_tokens: int = 200, temperature: float = 0.7):
    messages = [{"role": "user", "content": prompt}]
    formatted_prompt = tokenizer.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    inputs = tokenizer(formatted_prompt, return_tensors="pt")

    with torch.no_grad():
        output = model.generate(
            **inputs,
            max_new_tokens=int(max_tokens),
            temperature=temperature,
            do_sample=True,
            top_p=0.9,
            repetition_penalty=1.15,
            pad_token_id=tokenizer.eos_token_id,
        )

    response = tokenizer.decode(
        output[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True
    )
    return response

demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Slider(1, 512, value=200, label="Max tokens"),
        gr.Slider(0.1, 1.0, value=0.7, label="Temperature"),
    ],
    outputs="text",
    title="WINTER-FROST-2-PRO",
)
demo.launch()