import torch
import gradio as gr
import spaces
from transformers import pipeline

pipe = pipeline(
    "text-generation",
    model="Bouquets/StrikeGPT-R1-Zero-8B",
    dtype=torch.bfloat16,
    device_map="auto",
)

SYSTEM_PROMPT = (
    "You are a helpful AI assistant. "
    "Always provide complete, detailed and well-structured answers."
)

@spaces.GPU
def chat(message, history):
    # Gradio history is ignored here for simplicity and compatibility.
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": message},
    ]

    prompt = pipe.tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True,
    )

    output = pipe(
        prompt,
        max_new_tokens=2048,
        temperature=0.6,
        top_p=0.9,
        do_sample=True,
        repetition_penalty=1.1,
        return_full_text=False,
        pad_token_id=pipe.tokenizer.eos_token_id,
    )

    return output[0]["generated_text"].strip()


demo = gr.ChatInterface(
    fn=chat,
    title="🤖 StrikeGPT-R1 Assistant",
    description="Powered by Bouquets/StrikeGPT-R1-Zero-8B",
)

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