import gradio as gr
from transformers import pipeline

# Model name
MODEL_NAME = "jerteh/gpt2-vrabac"

# Initialize text generation pipeline
generator = pipeline(
    "text-generation",
    model=MODEL_NAME
)

def generate(text, temperature, top_p):
    if not text.strip():
        return "Molimo unesite neki tekst."

    output = generator(
        text,
        max_new_tokens=50,
        temperature=temperature,
        top_p=top_p,
        do_sample=True,
        no_repeat_ngram_size=3,
        repetition_penalty=1.2,
        num_return_sequences=1
    )

    return output[0]["generated_text"]

demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(
            label="Započnite rečenicu (srpski)",
            placeholder="Npr: Danas je lep dan..."
        ),
        gr.Slider(
            minimum=0.2,
            maximum=1.5,
            value=0.8,
            step=0.1,
            label="Temperatura (kreativnost)"
        ),
        gr.Slider(
            minimum=0.5,
            maximum=1.0,
            value=0.9,
            step=0.05,
            label="Top-p (raznovrsnost)"
        )
    ],
    outputs=gr.Textbox(label="AI nastavak"),
    title="Srpski AI Generator (Vrabac)",
    description="GPT-2 Vrabac sa kontrolom kreativnosti i bez ponavljanja 🐦"
)

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