import gradio as gr
from sentence_transformers import SentenceTransformer
from transformers import pipeline

print("⏳ Loading Models... (This happens once on startup)")

# 1. Load Embedding Model (MiniLM)
# Good for: Semantic search, matching influencers to brands
embedder = SentenceTransformer('all-MiniLM-L6-v2')

# 2. Load Sentiment Model (Twitter-RoBERTa)
# Good for: Checking if influencer content is positive/negative/toxic
sentiment_task = pipeline(
    "sentiment-analysis", 
    model="cardiffnlp/twitter-roberta-base-sentiment-latest", 
    tokenizer="cardiffnlp/twitter-roberta-base-sentiment-latest"
)

def get_embeddings(text):
    # Returns a list of 384 numbers (vector)
    vector = embedder.encode(text)
    return vector.tolist()

def get_sentiment(text):
    # Returns label (Positive/Negative) and score
    result = sentiment_task(text)[0]
    return {"label": result['label'], "score": float(result['score'])}

# --- UI & API DEFINITION ---
with gr.Blocks() as demo:
    gr.Markdown("# Flunzup NLP API")
    
    with gr.Tab("Sentiment Analysis"):
        txt_in = gr.Textbox(label="Text to Analyze")
        out_sent = gr.JSON(label="Sentiment Result")
        btn_sent = gr.Button("Analyze Sentiment")
        # EXPOSED API: /sentiment
        btn_sent.click(get_sentiment, inputs=txt_in, outputs=out_sent, api_name="sentiment")

    with gr.Tab("Embeddings"):
        txt_embed = gr.Textbox(label="Text to Embed")
        out_embed = gr.JSON(label="Vector Output")
        btn_embed = gr.Button("Generate Embedding")
        # EXPOSED API: /embedding
        btn_embed.click(get_embeddings, inputs=txt_embed, outputs=out_embed, api_name="embedding")

demo.launch()