import gradio as gr
from zipfile import ZipFile
import os


# Importing the sentiment_analysis function
def sentiment_analysis(dated_input):
    tokenizer = AutoTokenizer.from_pretrained("aliciiavs/sentiment-analysis-whatsapp2")
    model = AutoModelForSequenceClassification.from_pretrained("aliciiavs/sentiment-analysis-whatsapp2")

    # Tokenize input
    inputs = tokenizer(dated_input, padding=True, return_tensors="pt")

    # Forward pass through the model
    with torch.no_grad():
        outputs = model(**inputs)

    # Get predicted probabilities and predicted label
    probabilities = torch.softmax(outputs.logits, dim=1)
    predicted_label = torch.argmax(probabilities, dim=1)

    # Convert the predicted label tensor to a Python integer
    predicted_label = predicted_label.item()

    # Map predicted label index to sentiment label
    label_dic = {0: 'sadness', 1: 'joy', 2: 'love', 3: 'anger', 4: 'fear', 5: 'surprise'}

    # Return the predicted sentiment label instead of printing it
    return label_dic[predicted_label]

# Define a Gradio interface
def sentiment_analysis_interface(zip_file):
    # Extract text from zip file
    with ZipFile(zip_file) as archive:
        # Assuming each file in the zip contains text
        text = ""
        for filename in archive.namelist():
            with archive.open(filename) as file:
                text += file.read().decode("utf-8")
    
    # Perform sentiment analysis
    predicted_sentiment = sentiment_analysis(text)
    
    # Return the predicted sentiment label
    return f"Predicted sentiment label: {predicted_sentiment}"

# Create a Gradio interface
gr.Interface(
    fn=sentiment_analysis_interface,
    inputs=gr.File(type="filepath", label="Upload a zip file"),
    outputs="text"
).launch()