# import csv
import gradio as gr
import pandas as pd
from transformers import pipeline
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoConfig

# from datasets import load_dataset


# Load the model and define the sentiment classifier
MODEL = "LiYuan/amazon-review-sentiment-analysis"
tokenizer = AutoTokenizer.from_pretrained(MODEL)
config = AutoConfig.from_pretrained(MODEL)
model = AutoModelForSequenceClassification.from_pretrained(MODEL)
pipe = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer, config=config)


def classify_sentiment(sentences):
    """
    Classify the sentiment of each sentence
    """
    predictions = pipe(sentences)

    # Extract the predicted labels and confidence scores from the predictions
    labels = [prediction["label"] for prediction in predictions]
    confidences = [prediction["score"] for prediction in predictions]

    return labels, confidences


def classify_sentiment_from_csv(csv_file):
    """
    Read the CSV file and extract the list of sentences
    """
    df = pd.read_csv(csv_file.name, delimiter=",")
    sentences = df["sentence"].tolist()

    # Classify the sentiment of the sentences
    labels, confidences = classify_sentiment(sentences)
    df["confidences"] = confidences
    df["labels"] = labels
    return df


def main():
    """
    Define the gradio app
    """
    iface = gr.Interface(
        fn=classify_sentiment_from_csv,
        inputs=gr.File(),
        outputs=gr.Dataframe(),
        live=True,
        # capture_session=True,
        allow_flagging="never",
    )

    iface.launch(enable_queue=False)


# debug:
# labels, confidence = classify_sentiment_from_csv("./reviews.csv")
# print(labels)


# Run the gradio app
if __name__ == "__main__":
    main()