import gradio as gr
from transformers import pipeline

# --------- Load Models ---------
sentiment_model = pipeline("sentiment-analysis")
summarizer = pipeline("summarization")
qa_model = pipeline("question-answering")
translator = pipeline("translation_en_to_fr")
zero_shot = pipeline("zero-shot-classification")


# --------- Pipeline Functions ---------
def run_app(task, text, context="", labels=""):
    if task == "Sentiment Analysis":
        result = sentiment_model(text)[0]
        return f"{result['label']} (confidence: {round(result['score'], 3)})"

    elif task == "Text Summarization":
        result = summarizer(
            text,
            max_length=120,
            min_length=30,
            do_sample=False
        )[0]
        return result["summary_text"]

    elif task == "Question Answering":
        if not context:
            return "Please provide a context passage."
        result = qa_model(question=text, context=context)
        return result["answer"]

    elif task == "English → French Translation":
        result = translator(text)[0]
        return result["translation_text"]

    elif task == "Zero-Shot Classification":
            if not labels:
                return "Please enter candidate labels (comma-separated)."
            label_list = [x.strip() for x in labels.split(",")]
            result = zero_shot(text, candidate_labels=label_list)
            lines = [
                f"{label}: {round(score, 3)}"
                for label, score in zip(result["labels"], result["scores"])
            ]
            return "\n".join(lines)

    return "Select a valid task."


# --------- Gradio UI ---------
with gr.Blocks(title="Hugging Face AI Playground") as demo:
    gr.Markdown(
        "## 🤗 Hugging Face AI App\n"
        "Select a task and run inference using pretrained Transformer models."
    )

    task = gr.Dropdown(
        choices=[
            "Sentiment Analysis",
            "Text Summarization",
            "Question Answering",
            "English → French Translation",
            "Zero-Shot Classification",
        ],
        value="Sentiment Analysis",
        label="Choose a Task"
    )

    text = gr.Textbox(lines=4, label="Input Text")

    context = gr.Textbox(
        lines=4,
        label="Context (only for Question Answering)",
        visible=True
    )

    labels = gr.Textbox(
        label="Candidate Labels (comma-separated, for Zero-Shot Classification)"
    )

    output = gr.Textbox(label="Model Output")

    run_button = gr.Button("Run")

    run_button.click(
        fn=run_app,
        inputs=[task, text, context, labels],
        outputs=output
    )

demo.launch()