# app.py — Gradio chat app for your fine-tuned text-generation model
import gradio as gr
from transformers import pipeline
import torch

MODEL_ID = "samandar1105/text-generation"  # ← your model repo

print(f"Loading model: {MODEL_ID}")
generator = pipeline(
    "text-generation",
    model=MODEL_ID,
    device=0 if torch.cuda.is_available() else -1,
)
print("Model loaded successfully!")


def respond(message, history, max_new_tokens, temperature):
    messages = []
    for user_msg, bot_msg in history:
        messages.append({"role": "user", "content": user_msg})
        if bot_msg:
            messages.append({"role": "assistant", "content": bot_msg})
    messages.append({"role": "user", "content": message})

    output = generator(
        messages,
        max_new_tokens=max_new_tokens,
        temperature=temperature,
        do_sample=temperature > 0,
    )
    reply = output[0]["generated_text"][-1]["content"]
    return reply


with gr.Blocks(title="My Fine-Tuned LLM", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # My Fine-Tuned Text Generator
        Chat with a model fine-tuned on instruction-following data.
        """
    )

    with gr.Accordion("Generation settings", open=False):
        max_new_tokens = gr.Slider(16, 512, value=200, step=8, label="Max new tokens")
        temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.1, label="Temperature")

    gr.ChatInterface(
        fn=respond,
        additional_inputs=[max_new_tokens, temperature],
        examples=[
            ["Write a short story about a robot who learns to paint.", 200, 0.8],
            ["Explain quantum entanglement simply.", 200, 0.5],
            ["Give me 3 ideas for a weekend trip near the mountains.", 200, 0.7],
        ],
    )

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