Update app.py

app.py

@@ -99,7 +99,7 @@ class TextDetectionApp:
         with torch.no_grad():
             detection_score = self.ff_model(self.generate_ff_input(self.api_huggingface(text)))[0][0].item()
             # Return result based on the score threshold
-            return "Generated Text" if detection_score > 0.5 else "Human-Written"
+            return detection_score
 
     def classify_text(self, text, model_choice):
         """
@@ -110,7 +110,7 @@ class TextDetectionApp:
             model_choice (str): The model to use ('DeBERTa', 'RoBERTa', 'BERT', 'DistilBERT', or 'Feedforward').
 
         Returns:
-            str: The classification result.
+            str: The classification result including prediction scores.
         """
         if model_choice == 'DeBERTa':
             # Tokenize input
@@ -121,9 +121,12 @@ class TextDetectionApp:
 
             # Get classification results
             logits = outputs.logits
-            predicted_class_id = logits.argmax().item()
-            label = "Generated Text" if predicted_class_id == 1 else "Human-Written"
-            return f"{label} )"
+            scores = torch.softmax(logits, dim=1)[0]
+            generated_score = scores[1].item()
+            human_written_score = scores[0].item()
+            label = "Generated Text" if generated_score > 0.5 else "Human-Written"
+            return f"{label} ({generated_score * 100:.2f}% Generated, {human_written_score * 100:.2f}% Human-Written)"
+
         elif model_choice == 'RoBERTa':
             # Tokenize input
             inputs = self.roberta_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
@@ -133,9 +136,12 @@ class TextDetectionApp:
 
             # Get classification results
             logits = outputs.logits
-            predicted_class_id = logits.argmax().item()
-            label = "Generated Text" if predicted_class_id == 1 else "Human-Written"
-            return f"{label} )"
+            scores = torch.softmax(logits, dim=1)[0]
+            generated_score = scores[1].item()
+            human_written_score = scores[0].item()
+            label = "Generated Text" if generated_score > 0.5 else "Human-Written"
+            return f"{label} ({generated_score * 100:.2f}% Generated, {human_written_score * 100:.2f}% Human-Written)"
+
         elif model_choice == 'BERT':
             # Tokenize input
             inputs = self.bert_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
@@ -145,9 +151,12 @@ class TextDetectionApp:
 
             # Get classification results
             logits = outputs.logits
-            predicted_class_id = logits.argmax().item()
-            label = "Generated Text" if predicted_class_id == 1 else "Human-Written"
-            return f"{label} )"
+            scores = torch.softmax(logits, dim=1)[0]
+            generated_score = scores[1].item()
+            human_written_score = scores[0].item()
+            label = "Generated Text" if generated_score > 0.5 else "Human-Written"
+            return f"{label} ({generated_score * 100:.2f}% Generated, {human_written_score * 100:.2f}% Human-Written)"
+
         elif model_choice == 'DistilBERT':
             # Tokenize input
             inputs = self.distilbert_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
@@ -157,14 +166,19 @@ class TextDetectionApp:
 
             # Get classification results
             logits = outputs.logits
-            predicted_class_id = logits.argmax().item()
-            label = "Generated Text" if predicted_class_id == 1 else "Human-Written"
-            return f"{label} )"
+            scores = torch.softmax(logits, dim=1)[0]
+            generated_score = scores[1].item()
+            human_written_score = scores[0].item()
+            label = "Generated Text" if generated_score > 0.5 else "Human-Written"
+            return f"{label} ({generated_score * 100:.2f}% Generated, {human_written_score * 100:.2f}% Human-Written)"
 
-        elif model_choice == 'Feedforward':
+        elif model_choice == 'DAIGT-Model':
             # Run feedforward detection
-            detection_result = self.detect_text(text)
-            return f"{detection_result}"
+            detection_score = self.detect_text(text)
+            label = "Generated Text" if detection_score > 0.5 else "Human-Written"
+            generated_score = detection_score
+            human_written_score = 1 - detection_score
+            return f"{label} ({generated_score * 100:.2f}% Generated, {human_written_score * 100:.2f}% Human-Written)"
 
         else:
             return "Invalid model selection."
@@ -182,7 +196,7 @@ iface = gr.Interface(
     ],
     outputs="text",
     title="Text Classification with Multiple Models",
-    description="Classify text as generated or human-written using DeBERTa, RoBERTa, BERT, DistilBERT, or a custom Feedforward model."
+    description="Classify text as generated or human-written using DeBERTa, RoBERTa, BERT, DistilBERT, or a custom Feedforward model. See the confidence percentages for each prediction."
 )
 
 iface.launch()