import gradio as gr
from transformers import pipeline

# Initialize pipelines
sentiment_pipeline = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
classification_pipeline = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
ner_pipeline = pipeline("ner", model="dslim/bert-base-NER", aggregation_strategy="simple")
summarization_pipeline = pipeline("summarization", model="Curative/t5-summarizer-cnn")

# Define candidate labels for classification
candidate_labels = [
    "technology", "sports", "business", "politics",
    "health", "science", "travel", "entertainment"
]

def process(text, features):
    result = {}
    if "Summarization" in features:
        summary = summarization_pipeline(text, max_length=150, min_length=40, do_sample=False)
        result["summary"] = summary[0]["summary_text"]
    if "Sentiment" in features:
        sentiment = sentiment_pipeline(text)[0]
        result["sentiment"] = {"label": sentiment["label"], "score": sentiment["score"]}
    if "Classification" in features:
        classification = classification_pipeline(text, candidate_labels=candidate_labels)
        result["classification"] = dict(zip(classification["labels"], classification["scores"]))
    if "Entities" in features:
        entities = ner_pipeline(text)
        result["entities"] = [{"word": entity["word"], "type": entity["entity_group"]} for entity in entities]
    return result

# Build Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("## 🛠️ Multi-Feature NLP Service")
    inp = gr.Textbox(lines=6, placeholder="Enter your text here…")
    feats = gr.CheckboxGroup(
        ["Summarization", "Sentiment", "Classification", "Entities"],
        label="Select features to run"
    )
    btn = gr.Button("Run")
    out = gr.JSON(label="Results")
    btn.click(process, [inp, feats], out)

demo.queue(api_open=True).launch()