Update app.py

app.py

@@ -3,11 +3,6 @@ import pandas as pd
 import numpy as np
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer
 
-# load tokenizer and model, create trainer
-model_name = "j-hartmann/emotion-english-distilroberta-base"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSequenceClassification.from_pretrained(model_name)
-trainer = Trainer(model=model) 
 
 # summary function - test for single gradio function interfrace
 def bulk_function(filename):
@@ -22,11 +17,47 @@ def bulk_function(filename):
       def __getitem__(self, idx):
           return {k: v[idx] for k, v in self.tokenized_texts.items()}
 
-  # read file lines
-  with open(filename.name, "r") as f:
-    lines = f.readlines()
-  # expects unnamed:0 or index, col name -> strip both
-  lines_s = [item.split("\n")[0].split(",")[-1] for item in lines][1:]
+  # load tokenizer and model, create trainer
+  model_name = "j-hartmann/emotion-english-distilroberta-base"
+  tokenizer = AutoTokenizer.from_pretrained(model_name)
+  model = AutoModelForSequenceClassification.from_pretrained(model_name)
+  trainer = Trainer(model=model)  
+  print(filename, type(filename))
+  print(filename.name)
+
+
+  # check type of input file
+  if filename.name.split(".")[1] == "csv":
+    print("entered")
+    # read file, drop index if exists
+    df_input = pd.read_csv(filename.name, index_col=False)
+    if df_input.columns[0] == "Unnamed: 0":
+      df_input = df_input.drop("Unnamed: 0", axis=1)
+  elif filename.name.split(".")[1] == "xlsx":
+    df_input = pd.read_excel(filename.name, index_col=False)
+    # handle Unnamed
+    if df_input.columns[0] == "Unnamed: 0":
+      df_input = df_input.drop("Unnamed: 0", axis=1)
+  else:
+    return
+
+
+  # read csv
+  # even if index given, drop it
+  #df_input = pd.read_csv(filename.name, index_col=False)
+  #print("df_input", df_input)
+  
+  # expect csv format to be in: 
+  # 1: ID
+  # 2: Texts
+  # no index
+  # store ids in ordered list
+  ids = df_input[df_input.columns[0]].to_list()
+
+  # store sentences in ordered list
+  # expects sentences to be in second col
+  # of csv with two cols
+  lines_s = df_input[df_input.columns[1]].to_list()
 
     # Tokenize texts and create prediction data set
   tokenized_texts = tokenizer(lines_s,truncation=True,padding=True)
@@ -39,10 +70,13 @@ def bulk_function(filename):
   preds = predictions.predictions.argmax(-1)
   labels = pd.Series(preds).map(model.config.id2label)
   scores = (np.exp(predictions[0])/np.exp(predictions[0]).sum(-1,keepdims=True)).max(1)
+
+  # round scores
+  scores_rounded = [round(score, 2) for score in scores]
+
     # scores raw
   temp = (np.exp(predictions[0])/np.exp(predictions[0]).sum(-1,keepdims=True))
 
-    # work in progress
   # container
   anger = []
   disgust = []
@@ -54,24 +88,27 @@ def bulk_function(filename):
 
   # extract scores (as many entries as exist in pred_texts)
   for i in range(len(lines_s)):
-    anger.append(temp[i][0])
-    disgust.append(temp[i][1])
-    fear.append(temp[i][2])
-    joy.append(temp[i][3])
-    neutral.append(temp[i][4])
-    sadness.append(temp[i][5])
-    surprise.append(temp[i][6])
+    anger.append(round(temp[i][0], 3))
+    disgust.append(round(temp[i][1], 3))
+    fear.append(round(temp[i][2], 3))
+    joy.append(round(temp[i][3], 3))
+    neutral.append(round(temp[i][4], 3))
+    sadness.append(round(temp[i][5], 3))
+    surprise.append(round(temp[i][6], 3))
 
   # define df
-  df = pd.DataFrame(list(zip(lines_s,preds,labels,scores,  anger, disgust, fear, joy, neutral, sadness, surprise)), columns=['text','pred','label','score', 'anger', 'disgust', 'fear', 'joy', 'neutral', 'sadness', 'surprise'])
-
+  df = pd.DataFrame(list(zip(ids,lines_s,labels,scores_rounded, anger, disgust, fear, joy, neutral, sadness, surprise)), columns=[df_input.columns[0], df_input.columns[1],'label','score', 'anger', 'disgust', 'fear', 'joy', 'neutral', 'sadness', 'surprise'])
+  print(df)
   # save results to csv
   YOUR_FILENAME = filename.name.split(".")[0] + "_emotion_predictions" + ".csv"  # name your output file
-  df.to_csv(YOUR_FILENAME)
+  df.to_csv(YOUR_FILENAME, index=False)
 
   # return dataframe for space output
   return YOUR_FILENAME
   
-gr.Interface(bulk_function, [gr.inputs.File(file_count="single", type="file", label="csv", optional=False),],["file"],
-examples=[['emotion_examples.csv'],],
+gr.Interface(bulk_function, inputs=[gr.inputs.File(file_count="single", type="file", label="Upload file", optional=False),],
+             outputs=[gr.outputs.File(label="Output file")],
+            # examples=[["YOUR_FILENAME.csv"]], # computes, doesn't export df so far
+             theme="huggingface",
+             allow_flagging=False,
              ).launch(debug=True)