DivEye - PR (fix bugs, unmodularize)

app.py

@@ -18,95 +18,75 @@ import os
 
 theme = gr.Theme.from_hub("gstaff/xkcd") 
 
-class Diversity:
-    def __init__(self, model, tokenizer, device):
-        self.tokenizer = tokenizer
-        self.model = model
-        self.device = device
-        
-    def compute_log_likelihoods(self, text):
-        tokens = self.tokenizer.encode(text, return_tensors="pt", truncation=True, max_length=1024).to(self.device)
-        with torch.no_grad():
-            outputs = self.model(tokens, labels=tokens)
-        logits = outputs.logits
-        shift_logits = logits[:, :-1, :].squeeze(0)
-        shift_labels = tokens[:, 1:].squeeze(0)
-        log_probs = torch.log_softmax(shift_logits.float(), dim=-1)
-        token_log_likelihoods = log_probs[range(shift_labels.shape[0]), shift_labels].cpu().numpy()
-        return token_log_likelihoods
+# ===========================================================
+@spaces.GPU
+def evaluate(text):
+    global model, div_model, div_tokenizer, bi_model, bi_tokenizer
+
+    # =====================================================================
+    # DivEye features
+    diveye_features = []
+    # 1. Token log likelihoods
+    tokens = div_tokenizer.encode(text, return_tensors="pt", truncation=True, max_length=1024).to(div_model.device)
+    with torch.no_grad():
+        outputs = div_model(tokens, labels=tokens)
+    logits = outputs.logits
+    shift_logits = logits[:, :-1, :].squeeze(0)
+    shift_labels = tokens[:, 1:].squeeze(0)
+    log_probs = torch.log_softmax(shift_logits.float(), dim=-1)
+    token_log_likelihoods = log_probs[range(shift_labels.shape[0]), shift_labels].cpu().numpy()
     
-    def compute_surprisal(self, text):
-        log_likelihoods = self.compute_log_likelihoods(text)
-        return -log_likelihoods
-        
-    def compute_features(self, text):
-        surprisals = self.compute_surprisal(text)
-        log_likelihoods = self.compute_log_likelihoods(text)
-        if len(surprisals) < 10 or len(log_likelihoods) < 3:
-            return None
+    # 2. Surprisal
+    surprisals = -token_log_likelihoods
 
-        s = np.array(surprisals)
-        mean_s, std_s, var_s, skew_s, kurt_s = np.mean(s), np.std(s), np.var(s), skew(s), kurtosis(s)
-        diff_s = np.diff(s)
-        mean_diff, std_diff = np.mean(diff_s), np.std(diff_s)
-        first_order_diff = np.diff(log_likelihoods)
-        second_order_diff = np.diff(first_order_diff)
-        var_2nd = np.var(second_order_diff)
-        entropy_2nd = entropy(np.histogram(second_order_diff, bins=20, density=True)[0])
-        autocorr_2nd = np.corrcoef(second_order_diff[:-1], second_order_diff[1:])[0, 1] if len(second_order_diff) > 1 else 0
-        comp_ratio = len(zlib.compress(text.encode('utf-8'))) / len(text.encode('utf-8'))
+    if len(surprisals) < 10 or len(token_log_likelihoods) < 3:
+        diveye_features = [0.0] * 11
+    
+    s = np.array(surprisals)
+    mean_s, std_s, var_s, skew_s, kurt_s = np.mean(s), np.std(s), np.var(s), skew(s), kurtosis(s)
+    diff_s = np.diff(s)
+    mean_diff, std_diff = np.mean(diff_s), np.std(diff_s)
+    first_order_diff = np.diff(token_log_likelihoods)
+    second_order_diff = np.diff(first_order_diff)
+    var_2nd = np.var(second_order_diff)
+    entropy_2nd = entropy(np.histogram(second_order_diff, bins=20, density=True)[0])
+    autocorr_2nd = np.corrcoef(second_order_diff[:-1], second_order_diff[1:])[0, 1] if len(second_order_diff) > 1 else 0
+    comp_ratio = len(zlib.compress(text.encode('utf-8'))) / len(text.encode('utf-8'))
 
-        return [mean_s, std_s, var_s, skew_s, kurt_s, mean_diff, std_diff, var_2nd, entropy_2nd, autocorr_2nd, comp_ratio]
+    diveye_features = [mean_s, std_s, var_s, skew_s, kurt_s, mean_diff, std_diff, var_2nd, entropy_2nd, autocorr_2nd, comp_ratio]
+    # =====================================================================
 
-class BiScope:
-    def __init__(self, model, tokenizer, device):
-        self.COMPLETION_PROMPT_ONLY = "Complete the following text: "
-        self.tokenizer = tokenizer
-        self.model = model
-        self.device = device
-    
-    def compute_fce_loss(self, logits, targets, text_slice):
-        return CrossEntropyLoss(reduction='none')(
+    # =====================================================================
+    # BiScope features
+    COMPLETION_PROMPT_ONLY = "Complete the following text: "
+    prompt_ids = bi_tokenizer(COMPLETION_PROMPT_ONLY, return_tensors='pt').input_ids.to(bi_model.device)
+    text_ids = bi_tokenizer(text, return_tensors='pt', max_length=2000, truncation=True).input_ids.to(bi_model.device)
+    combined_ids = torch.cat([prompt_ids, text_ids], dim=1)
+    text_slice = slice(prompt_ids.shape[1], combined_ids.shape[1])
+
+    outputs = bi_model(input_ids=combined_ids)
+    logits = outputs.logits
+    targets = combined_ids[0][text_slice]
+
+    fce_loss = CrossEntropyLoss(reduction='none')(
             logits[0, text_slice.start-1:text_slice.stop-1, :],
             targets
         ).detach().cpu().numpy()
-    
-    def compute_bce_loss(self, logits, targets, text_slice):
-        return CrossEntropyLoss(reduction='none')(
+    bce_loss = CrossEntropyLoss(reduction='none')(
             logits[0, text_slice, :],
             targets
         ).detach().cpu().numpy()
     
-    def detect_single_sample(self, sample):
-        prompt_ids = self.tokenizer(self.COMPLETION_PROMPT_ONLY, return_tensors='pt').input_ids.to(self.device)
-        text_ids = self.tokenizer(sample, return_tensors='pt', max_length=2000, truncation=True).input_ids.to(self.device)
-        combined_ids = torch.cat([prompt_ids, text_ids], dim=1)
-        text_slice = slice(prompt_ids.shape[1], combined_ids.shape[1])
-
-        outputs = self.model(input_ids=combined_ids)
-        logits = outputs.logits
-        targets = combined_ids[0][text_slice]
-        
-        fce_loss = self.compute_fce_loss(logits, targets, text_slice)
-        bce_loss = self.compute_bce_loss(logits, targets, text_slice)
-
-        features = []
-        for p in range(1, 10):
-            split = len(fce_loss) * p // 10
-            fce_clipped = np.nan_to_num(np.clip(fce_loss[split:], -1e6, 1e6), nan=0.0, posinf=1e6, neginf=-1e6)
-            bce_clipped = np.nan_to_num(np.clip(bce_loss[split:], -1e6, 1e6), nan=0.0, posinf=1e6, neginf=-1e6)
-            features.extend([
-                np.mean(fce_clipped), np.max(fce_clipped), np.min(fce_clipped), np.std(fce_clipped),
-                np.mean(bce_clipped), np.max(bce_clipped), np.min(bce_clipped), np.std(bce_clipped)
-            ])
-        return features
-
-# ===========================================================
-@spaces.GPU
-def evaluate(diveye, biscope, text):
-    global model
-    diveye_features = diveye.compute_features(text)
-    biscope_features = biscope.detect_single_sample(text)
+    biscope_features = []
+    for p in range(1, 10):
+        split = len(fce_loss) * p // 10
+        fce_clipped = np.nan_to_num(np.clip(fce_loss[split:], -1e6, 1e6), nan=0.0, posinf=1e6, neginf=-1e6)
+        bce_clipped = np.nan_to_num(np.clip(bce_loss[split:], -1e6, 1e6), nan=0.0, posinf=1e6, neginf=-1e6)
+        biscope_features.extend([
+            np.mean(fce_clipped), np.max(fce_clipped), np.min(fce_clipped), np.std(fce_clipped),
+            np.mean(bce_clipped), np.max(bce_clipped), np.min(bce_clipped), np.std(bce_clipped)
+        ])
+    # =====================================================================    
 
     for f in biscope_features:
         diveye_features.append(f)
@@ -133,7 +113,7 @@ def detect_ai_text(text):
         )
 
     # Call software
-    ai_prob = evaluate(diveye, biscope, text)
+    ai_prob = evaluate(text)
     human_prob = 1 - ai_prob
 
     if ai_prob > 0.7:
@@ -178,9 +158,6 @@ if torch.cuda.is_available():
     model = xgb.XGBClassifier()
     model.load_model(model_path)
 
-    diveye = Diversity(div_model, div_tokenizer, div_model.device)
-    biscope = BiScope(bi_model, bi_tokenizer, bi_model.device)
-
 # Gradio app setup
 with gr.Blocks(title="DivEye") as demo:
     gr.HTML("""