Update detector.py

detector.py

@@ -5,11 +5,10 @@ import os
 import spaces
 
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
-torch.set_grad_enabled(False)  # Disable gradients globally
 
 class CustomDetector:
-    def __init__(self, model_name="tiiuae/falcon-rw-1b", max_length=256):
-        self.device = "cuda" if torch.cuda.is_available() else "cpu"  # Prefer GPU
+    def __init__(self, model_name="tiiuae/falcon-rw-1b", max_length=512):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.model_name = model_name
         self.max_length = max_length
         self.tokenizer = None
@@ -17,16 +16,13 @@ class CustomDetector:
 
     @spaces.GPU
     def load_model(self):
-        """Load model and tokenizer on GPU."""
+        """Load model and tokenizer on GPU when called."""
         try:
             if self.tokenizer is None:
                 self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
             if self.model is None:
-                self.model = AutoModelForCausalLM.from_pretrained(
-                    self.model_name,
-                    torch_dtype=torch.bfloat16,  # Use bfloat16 for GPU efficiency
-                    device_map="cuda"  # Auto-map to GPU
-                )
+                self.model = AutoModelForCausalLM.from_pretrained(self.model_name, torch_dtype=torch.float16)
+                self.model.to(self.device)
                 self.model.eval()
             if self.tokenizer.pad_token is None:
                 self.tokenizer.pad_token = self.tokenizer.eos_token
@@ -34,43 +30,37 @@ class CustomDetector:
             raise RuntimeError(f"Failed to load model {self.model_name}: {str(e)}")
 
     @spaces.GPU
-    def my_detector(self, texts: list[str], batch_size: int = 8) -> list[float]:
-        """Compute perplexity-based scores for texts."""
+    def my_detector(self, texts: list[str]) -> list[float]:
         if self.model is None or self.tokenizer is None:
             self.load_model()
 
         try:
-            all_scores = []
-            for i in range(0, len(texts), batch_size):
-                batch_texts = texts[i:i + batch_size]
+            with torch.no_grad():
                 tokenized = self.tokenizer(
-                    batch_texts,
+                    texts,
                     truncation=True,
-                    padding="longest" if len(batch_texts) > 1 else False,  # Dynamic padding
+                    padding=True,
                     max_length=self.max_length,
-                    return_tensors="pt"
+                    return_tensors="pt",
                 )
-                input_ids = tokenized["input_ids"].to(self.device)
-                attention_mask = tokenized["attention_mask"].to(self.device)
+                tokenized = {k: v.to(self.device) for k, v in tokenized.items()}
+                input_ids = tokenized["input_ids"]
+                attention_mask = tokenized["attention_mask"]
+
+                outputs = self.model(**tokenized)
+                logits = outputs.logits[:, :-1, :]
+                labels = tokenized["input_ids"][:, 1:]
 
-                with torch.inference_mode():  # Ensure no gradients
-                    outputs = self.model(input_ids, attention_mask=attention_mask)
-                    logits = outputs.logits[:, :-1, :]
-                    labels = input_ids[:, 1:]
+                log_probs = F.log_softmax(logits, dim=-1)
+                ll_per_token = log_probs.gather(2, labels.unsqueeze(-1)).squeeze(-1)
+                attention_mask = tokenized["attention_mask"][:, 1:]
+                ll_per_sample = (ll_per_token * attention_mask).sum(dim=-1) / attention_mask.sum(dim=1).clamp(min=1)
 
-                    # Compute perplexity: mean negative log-likelihood
-                    log_probs = F.log_softmax(logits, dim=-1)
-                    neg_log_likelihood = F.cross_entropy(
-                        logits.view(-1, logits.size(-1)),
-                        labels.view(-1),
-                        reduction="none"
-                    ).view(labels.size())
-                    attention_mask = attention_mask[:, 1:]
-                    perplexity = (neg_log_likelihood * attention_mask).sum(dim=-1) / attention_mask.sum(dim=-1).clamp(min=1)
-                    all_scores.extend(perplexity.tolist())
+                neg_entropy = (log_probs.exp() * log_probs)
+                entropy_per_sample = -(neg_entropy.sum(dim=-1) * attention_mask).sum(-1) / attention_mask.sum(dim=1).clamp(min=1)
 
-                torch.cuda.synchronize()  # Ensure GPU operations complete
+                scores = (abs(entropy_per_sample + ll_per_sample)).cpu().tolist()
 
-            return all_scores
+            return scores
         except Exception as e:
-            raise RuntimeError(f"Error computing score: {str(e)}")
+            raise RuntimeError(f"Error computing score: {str(e)}")