updated prediction finding

app.py

@@ -1,22 +1,22 @@
 import gradio as gr
 from transformers import pipeline
+import numpy as np
+
 
-# Load your models
-# Adjust these lines according to how your models are set up
 roberta_base_detector = pipeline("text-classification", model="Models/fine_tuned/roberta-base-openai-detector-model", tokenizer="Models/fine_tuned/roberta-base-openai-detector-tokenizer")
 chatgpt_lli_hc3_detector = pipeline("text-classification", model="Models/fine_tuned/chatgpt-detector-lli-hc3-model", tokenizer="Models/fine_tuned/chatgpt-detector-lli-hc3-tokenizer")
 chatgpt_roberta_detector = pipeline("text-classification", model="Models/fine_tuned/chatgpt-detector-roberta-model", tokenizer="Models/fine_tuned/chatgpt-detector-roberta-tokenizer")
 
 def classify_text(text):
     # Get predictions from each model
-    roberta_base_pred = roberta_base_detector(text)[0]['label']
-    chatgpt_lli_hc3_pred = chatgpt_lli_hc3_detector(text)[0]['label']
-    chatgpt_roberta_pred = chatgpt_roberta_detector(text)[0]['label']
+    roberta_base_pred = np.argmax(roberta_base_detector(text)[0]['label'].to('cpu'))
+    chatgpt_lli_hc3_pred = np.argmax(chatgpt_lli_hc3_detector(text)[0]['label'].to('cpu'))
+    chatgpt_roberta_pred = np.argmax(chatgpt_roberta_detector(text)[0]['label'].to('cpu'))
 
     # Count the votes for AI and Human
     votes = {"AI": 0, "Human": 0}
     for pred in [roberta_base_pred, chatgpt_lli_hc3_pred, chatgpt_roberta_pred]:
-        if pred == "AI":
+        if pred == 1:
             votes["AI"] += 1
         else:
             votes["Human"] += 1