fixaidetector

text_detector.py

@@ -2,7 +2,7 @@ import math
 import statistics
 import numpy as np
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForCausalLM
 from collections import Counter
 
 
@@ -17,6 +17,7 @@ class AITextDetector:
     def __init__(self, model_name="roberta-base-openai-detector", device=None):
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
+        self.lm_model = AutoModelForCausalLM.from_pretrained("gpt2")
 
         if device:
             self.device = device
@@ -24,19 +25,17 @@ class AITextDetector:
             self.device = "cuda" if torch.cuda.is_available() else "cpu"
 
         self.model.to(self.device)
+        self.lm_model.to(self.device)
         self.model.eval()
 
     def _compute_perplexity(self, text: str) -> float:
         """
         Approximate perplexity using NLL from model.
         """
-        encodings = self.tokenizer(text, return_tensors="pt", truncation=True)
-        input_ids = encodings.input_ids.to(self.device)
-
+        encodings = self.tokenizer(text, return_tensors="pt", truncation=True).to(self.device)
         with torch.no_grad():
-            outputs = self.model(input_ids, labels=input_ids)
+            outputs = self.lm_model(**encodings, labels=encodings.input_ids)
             loss = outputs.loss.item()
-
         return math.exp(loss)
 
     def _compute_burstiness(self, text: str) -> float: