import gradio as gr
from transformers import pipeline

# ---------- Load models ----------
sentiment = pipeline("sentiment-analysis")                # DistilBERT SST-2
classifier = pipeline("zero-shot-classification",
                      model="facebook/bart-large-mnli")   # Zero-shot

def analyze_email(subject, body):
    text = subject + "\n" + (body or "")

    # Sentiment
    s_res = sentiment(text)[0]
    s_label = s_res["label"]
    s_score = s_res["score"]

    # Zero-shot custom labels
    labels = ["engaging", "spammy", "informative", "boring", "urgent"]
    z_res = classifier(text, labels)
    z_scores = {l: f"{s:.2f}" for l, s in zip(z_res["labels"], z_res["scores"])}

    # ---------- format output ----------
    out  = f"### Sentiment\n**{s_label}**  (confidence {s_score:.2f})\n\n"
    out += "### Quality scores\n"
    for l, s in z_scores.items():
        out += f"- **{l}** : {s}\n"
    return out

demo = gr.Interface(
    fn     = analyze_email,
    inputs = [gr.Textbox(label="Subject line"),
              gr.Textbox(lines=6, label="Email body (optional)")],
    outputs = gr.Markdown(),
    title   = "Email Quality & Sentiment Analyzer",
    description = "Combines a sentiment pipeline + zero-shot classification"
)

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