Update app.py

app.py

@@ -86,11 +86,12 @@ class TextDetectionApp:
         Detects whether the input text is generated or human-written using the Feedforward model.
 
         Returns:
-            float: The detection result.
+            str: The detection result indicating if the text is generated or human-written.
         """
         with torch.no_grad():
-            self.output = self.ff_model(self.generate_ff_input(self.api_huggingface(text)))[0][0].item()
-        return self.output
+            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" if detection_score > 0.5 else "Human-Written"
 
     def classify_text(self, text, model_choice):
         """
@@ -113,7 +114,8 @@ class TextDetectionApp:
             # Get classification results
             logits = outputs.logits
             predicted_class_id = logits.argmax().item()
-            return f"DeBERTa Prediction: Class {predicted_class_id}"
+            label = "Generated" if predicted_class_id == 1 else "Human-Written"
+            return f"DeBERTa Prediction: {label} (Class {predicted_class_id})"
 
         elif model_choice == 'RoBERTa':
             # Tokenize input
@@ -125,12 +127,13 @@ class TextDetectionApp:
             # Get classification results
             logits = outputs.logits
             predicted_class_id = logits.argmax().item()
-            return f"RoBERTa Prediction: Class {predicted_class_id}"
+            label = "Generated" if predicted_class_id == 1 else "Human-Written"
+            return f"RoBERTa Prediction: {label} (Class {predicted_class_id})"
 
         elif model_choice == 'Feedforward':
             # Run feedforward detection
-            detection_score = self.detect_text(text)
-            return f"Feedforward Detection Score: {detection_score}"
+            detection_result = self.detect_text(text)
+            return f"Feedforward Detection: {detection_result}"
 
         else:
             return "Invalid model selection."
@@ -148,7 +151,7 @@ iface = gr.Interface(
     ],
     outputs="text",
     title="Text Classification with Multiple Models",
-    description="Classify text using DeBERTa, RoBERTa, or a custom Feedforward model."
+    description="Classify text as generated or human-written using DeBERTa, RoBERTa, or a custom Feedforward model."
 )
 
 iface.launch()