import gradio as gr
from transformers import pipeline

# 1️⃣ Lazy‑load your pipelines
summarizer = None
sentiment  = None
classifier = None
ner        = None

def get_summarizer():
    global summarizer
    if summarizer is None:
        summarizer = pipeline("summarization", model="Curative/t5-summarizer-cnn")
    return summarizer

def get_sentiment():
    global sentiment
    if sentiment is None:
        sentiment = pipeline("sentiment-analysis", model="DT12the/distilbert-sentiment-analysis")
    return sentiment

def get_classifier():
    global classifier
    if classifier is None:
        classifier = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")
    return classifier

def get_ner():
    global ner
    if ner is None:
        ner = pipeline("ner", model="dslim/bert-base-NER", aggregation_strategy="simple")
    return ner

# 2️⃣ Processing function
def process(text, features):
    """Run only the selected features on the input text."""
    results = {}
    if "Summarization" in features:
        summ = get_summarizer()(text, max_length=150, min_length=40, do_sample=False)[0]["summary_text"]
        results["summary"] = summ
    if "Sentiment" in features:
        sent = get_sentiment()(text)[0]
        results["sentiment"] = sent
    if "Classification" in features:
        cls = get_classifier()(text)[0]
        results["classification"] = cls
    if "Entities" in features:
        ents = get_ner()(text)
        # Format entities as list of dicts
        results["entities"] = [{"word": e["word"], "type": e["entity_group"]} for e in ents]
    return results

# 3️⃣ Build the Gradio Blocks UI
with gr.Blocks() as demo:
    gr.Markdown("## 📚 Multi‑Feature NLP Demo")
    text_input = gr.Textbox(lines=5, placeholder="Enter your text here…")
    feature_select = gr.CheckboxGroup(
        choices=["Summarization", "Sentiment", "Classification", "Entities"],
        label="Select features to run",
        info="You can pick one or more models to apply"
    )
    run_button = gr.Button("Run")
    output = gr.JSON(label="Results")

    run_button.click(
        fn=process,
        inputs=[text_input, feature_select],
        outputs=output
    )

# 4️⃣ Launch with API enabled
demo.queue(api_open=True).launch()