Update app.py

app.py

@@ -12,6 +12,14 @@ class TextDetectionApp:
         self.roberta_tokenizer = AutoTokenizer.from_pretrained("zeyadusf/roberta-DAIGT-kaggle")
         self.roberta_model = AutoModelForSequenceClassification.from_pretrained("zeyadusf/roberta-DAIGT-kaggle")
 
+        # Load BERT model and tokenizer
+        self.bert_tokenizer = AutoTokenizer.from_pretrained("zeyadusf/bert-DAIGT-MODELS")
+        self.bert_model = AutoModelForSequenceClassification.from_pretrained("zeyadusf/bert-DAIGT-MODELS")
+
+        # Load DistilBERT model and tokenizer
+        self.distilbert_tokenizer = AutoTokenizer.from_pretrained("zeyadusf/distilbert-DAIGT-MODELS")
+        self.distilbert_model = AutoModelForSequenceClassification.from_pretrained("zeyadusf/distilbert-DAIGT-MODELS")
+
         # Load Feedforward model
         self.ff_model = torch.jit.load("model_scripted.pt")
 
@@ -91,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" if detection_score > 0.5 else "Human-Written"
+            return "Generated Text" if detection_score > 0.5 else "Human-Written"
 
     def classify_text(self, text, model_choice):
         """
@@ -99,7 +107,7 @@ class TextDetectionApp:
 
         Args:
             text (str): The input text to classify.
-            model_choice (str): The model to use ('DeBERTa', 'RoBERTa', or 'Feedforward').
+            model_choice (str): The model to use ('DeBERTa', 'RoBERTa', 'BERT', 'DistilBERT', or 'Feedforward').
 
         Returns:
             str: The classification result.
@@ -114,9 +122,8 @@ class TextDetectionApp:
             # Get classification results
             logits = outputs.logits
             predicted_class_id = logits.argmax().item()
-            label = "Generated" if predicted_class_id == 1 else "Human-Written"
-            return f"DeBERTa Prediction: {label} (Class {predicted_class_id})"
-
+            label = "Generated Text" if predicted_class_id == 1 else "Human-Written"
+            return f"{label} )"
         elif model_choice == 'RoBERTa':
             # Tokenize input
             inputs = self.roberta_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
@@ -127,13 +134,37 @@ class TextDetectionApp:
             # Get classification results
             logits = outputs.logits
             predicted_class_id = logits.argmax().item()
-            label = "Generated" if predicted_class_id == 1 else "Human-Written"
-            return f"RoBERTa Prediction: {label} (Class {predicted_class_id})"
+            label = "Generated Text" if predicted_class_id == 1 else "Human-Written"
+            return f"{label} )"
+        elif model_choice == 'BERT':
+            # Tokenize input
+            inputs = self.bert_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+
+            # Run model
+            outputs = self.bert_model(**inputs)
+
+            # 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} )"
+        elif model_choice == 'DistilBERT':
+            # Tokenize input
+            inputs = self.distilbert_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+
+            # Run model
+            outputs = self.distilbert_model(**inputs)
+
+            # 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} )"
 
         elif model_choice == 'Feedforward':
             # Run feedforward detection
             detection_result = self.detect_text(text)
-            return f"Feedforward Detection: {detection_result}"
+            return f"{detection_result}"
 
         else:
             return "Invalid model selection."
@@ -147,11 +178,11 @@ iface = gr.Interface(
     fn=app.classify_text,
     inputs=[
         gr.Textbox(lines=2, placeholder="Enter your text here..."),
-        gr.Radio(choices=["DeBERTa", "RoBERTa", "Feedforward"], label="Model Choice")
+        gr.Radio(choices=["DeBERTa", "RoBERTa", "BERT", "DistilBERT", "Feedforward"], label="Model Choice")
     ],
     outputs="text",
     title="Text Classification with Multiple Models",
-    description="Classify text as generated or human-written using DeBERTa, RoBERTa, or a custom Feedforward model."
+    description="Classify text as generated or human-written using DeBERTa, RoBERTa, BERT, DistilBERT, or a custom Feedforward model."
 )
 
 iface.launch()