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

MODEL_ID = "tuf601121/scriptmodel"

device = 0 if torch.cuda.is_available() else -1

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForCausalLM.from_pretrained(MODEL_ID)

generator = pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    device=device
)

def generate_story(prompt, max_tokens, temperature, top_p):
    if not prompt.strip():
        return "⚠️ Kuch likhiye pehle."

    output = generator(
        prompt,
        max_new_tokens=max_tokens,
        do_sample=True,
        temperature=temperature,
        top_p=top_p,
        repetition_penalty=1.2,
        eos_token_id=tokenizer.eos_token_id
    )

    return output[0]["generated_text"]

with gr.Blocks(title="ScriptModel – Hindi Story AI") as demo:
    gr.Markdown("## ✍️ ScriptModel – Hindi Story Generator")
    gr.Markdown("Ye model aapki Hindi kahaniyon ke style me likhta hai.")

    prompt = gr.Textbox(
        label="✏️ Prompt",
        placeholder="एक गरीब किसान का बेटा शहर जाकर अपनी ज़िंदगी बदल देता है...",
        lines=4
    )

    with gr.Row():
        max_tokens = gr.Slider(50, 800, value=300, step=10, label="Length")
        temperature = gr.Slider(0.2, 1.2, value=0.8, step=0.05, label="Creativity")
        top_p = gr.Slider(0.5, 1.0, value=0.9, step=0.05, label="Top-p")

    generate_btn = gr.Button("🚀 Generate Story")
    output = gr.Textbox(label="📜 Output", lines=15)

    generate_btn.click(
        fn=generate_story,
        inputs=[prompt, max_tokens, temperature, top_p],
        outputs=output
    )

demo.launch()