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

model_id = "EleutherAI/gpt-neo-125M"

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

def generate_text(prompt, temperature, max_new_tokens):
    inputs = tokenizer(prompt, return_tensors="pt")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=int(max_new_tokens),
            do_sample=True,
            temperature=temperature,
            top_p=0.95,
            repetition_penalty=1.1,
            pad_token_id=tokenizer.eos_token_id
        )

    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    generated = full_text[len(prompt):]

    # 🔥 combine prompt + green continuation
    return f"""
    <div style="font-family:monospace; white-space:pre-wrap;">
        {prompt}<span style="color:#00ff88;">{generated}</span>
    </div>
    """

with gr.Blocks(css="""
textarea {font-family: monospace;}
""") as demo:

    gr.Markdown("# 🧠 TextPlayground")

    with gr.Row():
        with gr.Column(scale=3):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Type something like: Three reasons to start a succulent garden",
                lines=10
            )

            output = gr.HTML(label="Output")

            btn = gr.Button("Submit ⚡")

        with gr.Column(scale=1):
            temperature = gr.Slider(0.1, 1.5, value=0.7, label="Temperature")
            max_tokens = gr.Slider(10, 256, value=120, label="Max length")

    btn.click(
        fn=generate_text,
        inputs=[prompt, temperature, max_tokens],
        outputs=output
    )

demo.launch()