fixaidetector

text_detector.py

@@ -1,139 +1,165 @@
 import math
-import statistics
+import re
 import numpy as np
 import torch
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForCausalLM
+from transformers import (
+    AutoTokenizer,
+    AutoModelForSequenceClassification,
+    AutoModelForCausalLM,
+)
 from collections import Counter
 
 
 class AITextDetector:
     """
-    Advanced AI Text Detector (2025-ready):
+    AI Text Detector
+
     - Transformer classifier for AI vs Human
     - Metrics: perplexity, burstiness, repetition, semantic smoothness
-    - Category distribution (4-way classification for interpretability)
+    - Returns AI-vs-Human probability + category distribution
     """
 
     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
-        else:
-            self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        # Device setup
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
 
-        self.model.to(self.device)
-        self.lm_model.to(self.device)
+        # Classifier model & tokenizer
+        self.classifier_tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForSequenceClassification.from_pretrained(model_name).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).to(self.device)
+        # Language model for perplexity (lighter than full GPT-2 if needed)
+        self.lm_tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        self.lm_model = AutoModelForCausalLM.from_pretrained("gpt2").to(self.device)
+        self.lm_model.eval()
+
+    # ------------------ Metrics ------------------
+    def _compute_perplexity(self, text: str, max_length: int = 512):
+        """Compute perplexity using GPT-2 LM."""
+        encodings = self.lm_tokenizer(
+            text,
+            return_tensors="pt",
+            truncation=True,
+            max_length=max_length,
+        ).to(self.device)
+
         with torch.no_grad():
             outputs = self.lm_model(**encodings, labels=encodings.input_ids)
             loss = outputs.loss.item()
-        return math.exp(loss)
 
-    def _compute_burstiness(self, text: str) -> float:
-        """
-        Burstiness = variance / mean of sentence lengths.
-        """
-        sentences = [s.strip() for s in text.split(".") if s.strip()]
-        lengths = [len(s.split()) for s in sentences]
+        # Clamp to avoid overflow
+        return float(min(math.exp(loss), 1e4))
 
-        if len(lengths) < 2:
+    def _compute_burstiness(self, text: str):
+        """Variance of sentence lengths (burstiness)."""
+        sentences = [s.strip() for s in re.split(r"[.!?]", text) if s.strip()]
+        if len(sentences) < 2:
             return 0.0
+        lengths = [len(s.split()) for s in sentences]
+        return float(np.var(lengths))
 
-        return statistics.pvariance(lengths) / (np.mean(lengths) + 1e-8)
-
-    def _compute_repetition_score(self, text: str) -> float:
-        """
-        Measures how often words repeat.
-        High repetition = more likely AI.
-        """
-        words = [w.lower() for w in text.split() if w.isalpha()]
+    def _compute_repetition_score(self, text: str):
+        """Repetition = proportion of duplicate words."""
+        words = [w.lower() for w in re.findall(r"\b\w+\b", text)]
         if not words:
             return 0.0
+        counts = Counter(words)
+        repeated = sum(c - 1 for c in counts.values() if c > 1)
+        return repeated / len(words)
 
-        word_counts = Counter(words)
-        repetition = sum(c - 1 for c in word_counts.values() if c > 1) / len(words)
-        return repetition
-
-    def _compute_semantic_smoothness(self, text: str) -> float:
+    def _compute_semantic_smoothness(self, text: str):
         """
-        Semantic smoothness = similarity between consecutive sentences.
-        Higher = more consistent flow (AI often too smooth).
+        Semantic smoothness = avg cosine similarity between consecutive sentence embeddings.
+        Uses last hidden states instead of raw embeddings.
         """
-        sentences = [s.strip() for s in text.split(".") if s.strip()]
+        sentences = [s.strip() for s in re.split(r"[.!?]", text) if s.strip()]
         if len(sentences) < 2:
             return 1.0
 
-        embeddings = self.model.base_model.get_input_embeddings()(
-            torch.tensor([self.tokenizer.encode(s, truncation=True, max_length=32) for s in sentences]).to(self.device)
-        )
-        embeddings = embeddings.mean(dim=1).detach().cpu().numpy()
-
-        sims = []
+        embeddings = []
+        for s in sentences:
+            encodings = self.classifier_tokenizer(
+                s,
+                return_tensors="pt",
+                truncation=True,
+                padding=True,
+                max_length=128,
+            ).to(self.device)
+
+            with torch.no_grad():
+                outputs = self.model(
+                    **encodings,
+                    output_hidden_states=True,
+                )
+                hidden_states = outputs.hidden_states[-1]  # last layer
+                sent_emb = hidden_states.mean(dim=1).cpu().numpy()
+            embeddings.append(sent_emb)
+
+        similarities = []
         for i in range(len(embeddings) - 1):
-            v1, v2 = embeddings[i], embeddings[i + 1]
-            cos = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2) + 1e-8)
-            sims.append(cos)
+            a, b = embeddings[i], embeddings[i + 1]
+            num = float(np.dot(a, b.T))
+            denom = np.linalg.norm(a) * np.linalg.norm(b)
+            if denom > 0:
+                similarities.append(num / denom)
+        return float(np.mean(similarities)) if similarities else 1.0
+
+    # ------------------ Main detection ------------------
+    def detect(self, text: str):
+        """Run detection pipeline and return results."""
+        # Empty text case
+        if not text.strip():
+            return {
+                "ai_probability": 0.0,
+                "metrics": {},
+                "distribution": {},
+                "final_label": "empty",
+            }
+
+        # Classifier prediction
+        inputs = self.classifier_tokenizer(
+            text,
+            return_tensors="pt",
+            truncation=True,
+            padding=True,
+            max_length=512,
+        ).to(self.device)
 
-        return float(np.mean(sims))
-
-    def detect(self, text: str) -> dict:
-        """
-        Run detection and return structured JSON report.
-        """
-
-        # Model classification
-        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(self.device)
         with torch.no_grad():
             logits = self.model(**inputs).logits
-            probs = torch.softmax(logits, dim=-1).cpu().numpy()[0]
+            probs = torch.softmax(logits, dim=1).cpu().numpy()[0]
 
-        ai_prob = float(probs[1]) if len(probs) > 1 else 0.5
+        human_prob, ai_prob = float(probs[0]), float(probs[1])
 
-        # Compute metrics
+        # Extra metrics
         perplexity = self._compute_perplexity(text)
         burstiness = self._compute_burstiness(text)
         repetition = self._compute_repetition_score(text)
         smoothness = self._compute_semantic_smoothness(text)
 
-        # Create 4-category distribution (mock scaling from ai_prob + heuristics)
+        # Normalize distribution
         distribution = {
-            "AI-generated": round(ai_prob * 100 * (1 - repetition), 1),
-            "AI-generated & AI-refined": round(ai_prob * 100 * repetition, 1),
-            "Human-written & AI-refined": round((1 - ai_prob) * 100 * smoothness, 1),
-            "Human-written": round((1 - ai_prob) * 100 * (1 - smoothness), 1)
+            "Human-written": round(human_prob * 100, 2),
+            "AI-generated": round(ai_prob * 100 * (1 - repetition), 2),
+            "AI-generated & AI-refined": round(ai_prob * 100 * repetition, 2),
+            "Mixed": round(ai_prob * 100 * (1 - smoothness), 2),
         }
-
-        # Normalize so they sum to 100
         total = sum(distribution.values())
         if total > 0:
-            distribution = {k: round(v / total * 100, 1) for k, v in distribution.items()}
+            for k in distribution:
+                distribution[k] = round(distribution[k] / total * 100, 2)
 
-        overall_ai_probability = round(ai_prob, 2)
-        summary = f"{distribution['AI-generated']}% of text is likely AI"
+        # Final label
+        final_label = max(distribution, key=distribution.get)
 
         return {
-            "summary": summary,
-            "overall_ai_probability": overall_ai_probability,
-            "category_distribution": distribution,
+            "ai_probability": round(ai_prob, 4),
             "metrics": {
-                "perplexity": round(perplexity, 2),
+                "perplexity": round(perplexity, 3),
                 "burstiness": round(burstiness, 3),
-                "repetition_score": round(repetition, 3),
+                "repetition": round(repetition, 3),
                 "semantic_smoothness": round(smoothness, 3),
-                "ai_probability": overall_ai_probability
             },
-            "interpretation": (
-                "This detector uses structural patterns (perplexity, burstiness, repetition, semantic smoothness) "
-                "to estimate the likelihood of AI authorship. Results are probabilistic, not definitive. "
-                "Always apply judgment."
-            )
+            "distribution": distribution,
+            "final_label": final_label,
         }