import gradio as gr
from transformers import pipeline
import pandas as pd
import matplotlib.pyplot as plt
import tempfile

# Initialize the zero-shot classification pipeline
classifier = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")



# Define the classification function
def classify_text(document, labels):
    candidate_labels = labels.split(", ")
    res = classifier(document, candidate_labels=candidate_labels, multi_label=False)
    df = pd.DataFrame(res)
    
    # Create a temporary file to save the plot
    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile:
        df.plot.bar(x='labels', y='scores', legend=False)
        plt.title("Classification Results")
        plt.xlabel("Labels")
        plt.ylabel("Scores")
        plt.tight_layout()
        plt.savefig(tmpfile.name)
        plt.close()
    
    return df, tmpfile.name

# Define the example inputs and outputs
examples = [
    ["It was about eleven o’clock in the morning, mid October, with the sun not shining and a look of hard wet rain in the clearness of the foothills. I was wearing my powder-blue suit, with dark blue shirt, tie and display handkerchief, black brogues, black wool socks with dark blue clocks on them. I was neat, clean, shaved and sober, and I didn’t care who knew it. I was everything the well-dressed private detective ought to be. I was calling on four million dollars.",
     "history, crime, fantasy"],
]

# Create Gradio interface
interface = gr.Interface(
    fn=classify_text,
    inputs=[
        gr.Textbox(lines=10, label="Document"),
        gr.Textbox(lines=1, label="Candidate Labels (comma-separated)")
    ],
    outputs=[
        gr.Dataframe(type ="pandas",label="Classification Scores"),
        gr.Image(type="numpy", label="Classification Bar Plot")
    ],
    title="Text Genre Classification",
    description="Classify text into specified labels using zero-shot classification. Provide a document and candidate labels separated by commas.",
    examples=examples
)

# Launch the Gradio app
interface.launch(debug=False)