import gradio as gr
from transformers import pipeline

# pipeline as high level
pipe = pipeline("text-generation", model="HuggingFaceTB/SmolLM2-360M-Instruct", device="cpu")

# 3 adjectives + don't repeat + don't elaborate
system_instruction = """
You are given a short description of a person. Respond only with three different adjectives that best describe the person, based on the description. No explanation.
"""

def get_vibe(caption):
    if not caption or not caption.strip():
        raise gr.Error("No caption provided.")

    messages = [
        {"role": "system", "content": system_instruction},
        {"role": "user", "content": caption}
    ]
    output = pipe(messages, max_new_tokens=32, do_sample=True, temperature=0.8, top_p=0.9)[0]["generated_text"]

    # extract assistant reply
    assistant_text = next(
        (entry['content'] for entry in output if entry['role'] == 'assistant'),
        ""
    )

    # normalize to "word, word, word"
    cleaned = assistant_text.replace(", and", ",").replace(" and", ",")
    adjectives = [word.strip() for word in cleaned.split(",") if word.strip()]
    return ", ".join(adjectives[:3])

# api w/ gradio
api = gr.Interface(
    fn=get_vibe,
    inputs=gr.Textbox(label="Input Caption"),
    outputs="text"
)

api.launch(show_api=True)