Upload app.py

app.py

@@ -12,11 +12,11 @@ import tqdm
 nltk.download('punkt')
 
 # Define the device (GPU or CPU)
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 
 # Define the model and tokenizer
 checkpoint = "ieq/IEQ-BERT"
-tokenizer = AutoTokenizer.from_pretrained(checkpoint).to(device)
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
 model = AutoModelForSequenceClassification.from_pretrained(checkpoint).to(device)
 
 
@@ -212,6 +212,18 @@ def predict_from_csv(file, column_name, progress=gr.Progress()):
     labels_predicted = []
     prediction_scores = []
 
+    # Initialize empty lists for IEQ labels and scores
+    ieq1 = []
+    ieq2 = []
+    ieq3 = []
+    ieq4 = []
+    ieq5 = []
+    score1 = []
+    score2 = []
+    score3 = []
+    score4 = []
+    score5 = []
+
     # Preprocess text and make predictions
     for text_input in progress.tqdm(text_list, desc="Analysing data"):
         # Sleep to avoid rate limiting
@@ -248,9 +260,69 @@ def predict_from_csv(file, column_name, progress=gr.Progress()):
         labels_predicted.append(predicted_labels)
         prediction_scores.append(prediction_score)
 
-        # Append the predictions to the DataFrame
+        # Append to ieq1 to ieq5
+        for i in range(5):
+            if i < len(predicted_labels):
+                if i == 0:
+                    ieq1.append(predicted_labels[i])
+                elif i == 1:
+                    ieq2.append(predicted_labels[i])
+                elif i == 2:
+                    ieq3.append(predicted_labels[i])
+                elif i == 3:
+                    ieq4.append(predicted_labels[i])
+                elif i == 4:
+                    ieq5.append(predicted_labels[i])
+            else:
+                if i == 0:
+                    ieq1.append("-")
+                elif i == 1:
+                    ieq2.append("-")
+                elif i == 2:
+                    ieq3.append("-")
+                elif i == 3:
+                    ieq4.append("-")
+                elif i == 4:
+                    ieq5.append("-")
+
+        # Append to score1 to score5
+        for i in range(5):
+            if i < len(prediction_score):
+                if i == 0:
+                    score1.append(prediction_score[i])
+                elif i == 1:
+                    score2.append(prediction_score[i])
+                elif i == 2:
+                    score3.append(prediction_score[i])
+                elif i == 3:
+                    score4.append(prediction_score[i])
+                elif i == 4:
+                    score5.append(prediction_score[i])
+            else:
+                if i == 0:
+                    score1.append("-")
+                elif i == 1:
+                    score2.append("-")
+                elif i == 2:
+                    score3.append("-")
+                elif i == 3:
+                    score4.append("-")
+                elif i == 4:
+                    score5.append("-")
+
+    # Append the predictions to the DataFrame
     df_docs['IEQ_predicted'] = labels_predicted
     df_docs['prediction_scores'] = prediction_scores
+    df_docs['IEQ1'] = ieq1
+    df_docs['IEQ2'] = ieq2
+    df_docs['IEQ3'] = ieq3
+    df_docs['IEQ4'] = ieq4
+    df_docs['IEQ5'] = ieq5
+    df_docs['Score1'] = score1
+    df_docs['Score2'] = score2
+    df_docs['Score3'] = score3
+    df_docs['Score4'] = score4
+    df_docs['Score5'] = score5
 
     # Save the predictions to a CSV file
     df_docs.to_csv('IEQ_predictions.csv')