import gradio as gr
from transformers import pipeline, AutoTokenizer

model_names = [
    "distilgpt2",
    "gpt2",
]
 
# --- minimal caching so the model isn't reloaded every click ---
_pipe_cache = {}

def _get_pipe(model_name):
    if model_name not in _pipe_cache:
        tok = AutoTokenizer.from_pretrained(model_name)
        _pipe_cache[model_name] = pipeline(
            "text-generation",
            model=model_name,
            tokenizer=tok,
            device_map="auto",
            torch_dtype="auto",
        )
    return _pipe_cache[model_name]
# ---------------------------------------------------------------

def generate_with_choice(prompt, model_name):
    pipe = _get_pipe(model_name)
    out = pipe(
        prompt,
        max_new_tokens=50,
        do_sample=True,
        return_full_text=False
    )
    return out[0]["generated_text"]

demo2 = gr.Interface(
    fn=generate_with_choice,
    inputs=[
        gr.Textbox(lines=4, label="Enter Prompt"),
        gr.Dropdown(model_names, label="Choose Model"),
    ],
    outputs=gr.Textbox(lines=5, label="Output"),
    flagging_mode="never",
    title="Model Chooser Demo",
    description="Pick a model and generate text on the fly!",
    theme="soft",
)

demo2.queue().launch(share=True)