716this's picture
apply black formatting and add doc strings
6838471
raw
history blame contribute delete
No virus
1.79 kB
# 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()