#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

# Define the gradio app
def main():
  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()