from transformers import pipeline
import gradio as gr

# Load model
generator = pipeline(
    "text-generation",
    model="Qwen/Qwen2.5-1.5B-Instruct"
)

def define_word(word):
    word = word.strip()

    if not word:
        return "Please enter a word."

    prompt = (
        "Define the following word in one clear sentence. "
        "Do not add examples, commentary, or extra information.\n\n"
        f"WORD: {word}\n"
        "DEFINITION:"
    )

    out = generator(
        prompt,
        max_new_tokens=40,
        do_sample=False,
        temperature=0.0,
        repetition_penalty=1.1,
        pad_token_id=50256
    )

    result = out[0]["generated_text"][len(prompt):].strip()

    # Simple post-processing safety
    if not result:
        return "No definition available."

    return result


app = gr.Interface(
    fn=define_word,
    inputs=gr.Textbox(
        lines=1,
        placeholder="Enter a word (e.g. audit, risk, encryption)"
    ),
    outputs=gr.Textbox(
        lines=3,
        label="Definition"
    ),
    title="AI Dictionary (SLM Demo)",
    description=(
        "A controlled language model acting as a dictionary. "
        "One word in — one definition out."
    ),
    allow_flagging="never"
)

app.launch()