Update ai_text_detector_valid_final.py

ai_text_detector_valid_final.py

@@ -12,15 +12,31 @@ MODELS = {
     
 def load_model(model_id):
     tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForSequenceClassification.from_pretrained(model_id)
+    # Use the zero-shot classification pipeline for NLI models
+    if model_id == "roberta-large-mnli":
+        model = pipeline("zero-shot-classification", model=model_id, device=0 if torch.cuda.is_available() else -1)
+    else:
+        model = AutoModelForSequenceClassification.from_pretrained(model_id)
     return tokenizer, model
 
 def predict(text, tokenizer, model):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
-    with torch.no_grad():
-        outputs = model(**inputs)
-        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
-    return probs[0].numpy()  # [human_prob, ai_prob]
+    if isinstance(model, pipeline):
+        # Use the roberta-mnli model for zero-shot classification
+        candidate_labels = ["This text was written by a human.", "This text was written by an AI."]
+        result = model(text, candidate_labels)
+        
+        # The entailment score for each label is the probability
+        human_prob = result["scores"][result["labels"].index("This text was written by a human.")]
+        ai_prob = result["scores"][result["labels"].index("This text was written by an AI.")]
+        
+        return np.array([human_prob, ai_prob])
+    else:
+        # The existing code for other models
+        inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+        with torch.no_grad():
+            outputs = model(**inputs)
+            probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+        return probs[0].numpy()
 
 
 def verdict(ai_prob):