Create model.py

model.py

@@ -0,0 +1,199 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModelForCausalLM, T5EncoderModel, RobertaTokenizer
+from huggingface_hub import hf_hub_download
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+MAX_LEN = 256
+THRESHOLD = 0.475
+REPO_ID = "santh-cpu/ai_code_detect"
+
+class PolyglotMetricEngine(nn.Module):
+    def __init__(self, base_model):
+        super().__init__()
+        self.model = base_model
+
+    @torch.no_grad()
+    def forward(self, input_ids, attention_mask):
+        B, L = input_ids.shape
+        with torch.autocast(device_type='cuda', dtype=torch.float16):
+            logits_raw = self.model(input_ids=input_ids, attention_mask=attention_mask, use_cache=False).logits
+        logits_raw = logits_raw.float()
+        shift_logits = logits_raw[:, :-1, :].contiguous()
+        shift_labels = input_ids[:, 1:].contiguous()
+        shift_mask   = attention_mask[:, 1:].float()
+        log_probs_all = F.log_softmax(shift_logits, dim=-1)
+        probs_all     = log_probs_all.exp()
+        log_prob = log_probs_all.gather(2, shift_labels.unsqueeze(-1)).squeeze(-1)
+        K_max = min(1001, shift_logits.size(-1))
+        topk_vals, topk_idx = torch.topk(shift_logits, K_max, dim=-1)
+        rank_approx   = (log_probs_all.gather(2, topk_idx) > log_prob.unsqueeze(-1)).sum(-1).float() + 1.0
+        true_rank_log = torch.log1p(rank_approx)
+        top10_mass = log_probs_all.gather(2, topk_idx[:, :, :10]).exp().sum(dim=-1)
+        lp_topk    = log_probs_all.gather(2, topk_idx[:, :, :2])
+        gap_1_2    = (lp_topk[:, :, 0] - lp_topk[:, :, 1]).clamp(-20, 20)
+        entropy    = -(probs_all * log_probs_all).sum(dim=-1)
+        varentropy = (probs_all * (-log_probs_all - entropy.unsqueeze(-1))**2).sum(dim=-1)
+        r10_flag   = (true_rank_log <= math.log1p(10)).float()
+        r100_flag  = ((true_rank_log > math.log1p(10)) & (true_rank_log <= math.log1p(100))).float()
+        r1k_flag   = ((true_rank_log > math.log1p(100)) & (true_rank_log <= math.log1p(1000))).float()
+        rtail_flag = (true_rank_log > math.log1p(1000)).float()
+        valid_n    = shift_mask.sum(dim=1, keepdim=True).clamp(min=1)
+        lp_mean    = (log_prob * shift_mask).sum(1, keepdim=True) / valid_n
+        lp_var     = ((log_prob - lp_mean)**2 * shift_mask).sum(1, keepdim=True) / valid_n
+        lp_std     = lp_var.sqrt().clamp(min=1e-4)
+        surprisal_z   = ((log_prob - lp_mean) / lp_std) * shift_mask
+        entropy_shift = F.pad(entropy[:, :-1], (1, 0), value=0.)
+        entropy_delta = (entropy - entropy_shift) * shift_mask
+        cum_positions = torch.arange(1, L, device=input_ids.device).unsqueeze(0).float()
+        cum_rank      = (true_rank_log * shift_mask).cumsum(dim=1) / cum_positions
+        is_special    = torch.zeros_like(shift_mask)
+        m = shift_mask
+        token_feats_12 = torch.stack([
+            log_prob*m, true_rank_log*m, entropy*m, varentropy*m,
+            top10_mass*m, gap_1_2*m, surprisal_z, entropy_delta,
+            cum_rank*m, is_special*m, r10_flag*m, r100_flag*m,
+        ], dim=-1)
+        out_token = torch.zeros(B, MAX_LEN, 12, device=input_ids.device)
+        out_token[:, :L-1, :] = token_feats_12
+        seq_feats = self._compute_seq_feats(
+            log_prob, entropy, varentropy, top10_mass,
+            gap_1_2, surprisal_z, r10_flag, r100_flag,
+            r1k_flag, rtail_flag, shift_mask, valid_n
+        )
+        return out_token.detach(), seq_feats.detach()
+
+    def _compute_seq_feats(self, log_prob, entropy, varentropy, top10_mass,
+                           gap_1_2, surprisal_z, r10_flag, r100_flag,
+                           r1k_flag, rtail_flag, mask, valid_n):
+        feats = []
+        def masked_moments(x):
+            n    = valid_n.squeeze(1)
+            mu   = (x * mask).sum(1) / n
+            dev  = (x - mu.unsqueeze(1)) * mask
+            var  = (dev**2).sum(1) / n
+            std  = var.sqrt().clamp(min=1e-6)
+            skew = (dev**3).sum(1) / (n * std**3 + 1e-8)
+            kurt = (dev**4).sum(1) / (n * var**2 + 1e-8)
+            return mu, std, skew.clamp(-10, 10), kurt.clamp(0, 50)
+        for feat in [log_prob, entropy, varentropy, top10_mass]:
+            feats += list(masked_moments(feat))
+        for lag in [1, 5]:
+            e_shift      = F.pad(entropy[:, lag:], (0, lag)) * mask
+            e_norm       = entropy - (entropy*mask).sum(1, keepdim=True)/valid_n
+            e_shift_norm = e_shift - (e_shift*mask).sum(1, keepdim=True)/valid_n
+            num   = (e_norm * e_shift_norm * mask).sum(1)
+            denom = (((e_norm**2*mask).sum(1)+1e-8).sqrt() * ((e_shift_norm**2*mask).sum(1)+1e-8).sqrt())
+            feats.append((num/denom).clamp(-1, 1))
+        n = valid_n.squeeze(1)
+        for flag in [r10_flag, r100_flag, r1k_flag, rtail_flag]:
+            feats.append((flag*mask).sum(1)/n)
+        feats += list(masked_moments(gap_1_2)[:2])
+        feats += list(masked_moments(surprisal_z)[:2])
+        pos     = torch.arange(entropy.shape[1], device=entropy.device).float().unsqueeze(0)
+        pos_mu  = (pos*mask).sum(1)/n
+        ent_mu  = (entropy*mask).sum(1)/n
+        cov     = ((pos - pos_mu.unsqueeze(1)) * (entropy - ent_mu.unsqueeze(1)) * mask).sum(1)
+        var_pos = ((pos - pos_mu.unsqueeze(1))**2 * mask).sum(1)
+        feats.append((cov/(var_pos+1e-8)).clamp(-5, 5))
+        abs_surp = surprisal_z.abs()
+        mu_b, std_b, _, _ = masked_moments(abs_surp)
+        feats.append((std_b/(mu_b+1e-6)).clamp(0, 20))
+        feats.append((top10_mass*mask).sum(1)/n)
+        feats.append((((top10_mass-(top10_mass*mask).sum(1,keepdim=True)/valid_n)**2*mask).sum(1)/n).sqrt())
+        ent_median = entropy.median(dim=1).values.unsqueeze(1)
+        feats.append(((entropy < ent_median)*mask).sum(1)/n)
+        lp_std  = masked_moments(log_prob)[1]
+        ent_mu2 = (entropy*mask).sum(1)/n
+        feats.append((lp_std/(ent_mu2+1e-4)).clamp(0, 20))
+        return torch.nan_to_num(torch.stack(feats, dim=1), nan=0., posinf=20., neginf=-20.)
+
+class GatedTemporalMixer(nn.Module):
+    def __init__(self, dim, kernel_size=7):
+        super().__init__()
+        self.conv = nn.Conv1d(dim, dim*2, kernel_size, padding=(kernel_size-1), groups=dim)
+        self.norm = nn.LayerNorm(dim)
+    def forward(self, x):
+        h = self.conv(x.transpose(1,2))[:, :, :x.shape[1]]
+        gate, val = h.chunk(2, dim=1)
+        return self.norm((torch.sigmoid(gate)*val).transpose(1,2) + x)
+
+class PerTokenEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.feat_norm = nn.LayerNorm(12)
+        self.proj_in   = nn.Linear(12, 128)
+        self.mixer1    = GatedTemporalMixer(128, 7)
+        self.mixer2    = GatedTemporalMixer(128, 15)
+        self.ff = nn.Sequential(
+            nn.Linear(128, 256), nn.GELU(), nn.Dropout(0.1),
+            nn.Linear(256, 128), nn.LayerNorm(128)
+        )
+        self.attn_q   = nn.Linear(128, 1, bias=False)
+        self.proj_out = nn.Linear(128, 256)
+    def forward(self, x, mask):
+        x_proj = F.gelu(self.proj_in(self.feat_norm(x)))
+        mixed  = self.mixer2(self.mixer1(x_proj))
+        hidden = mixed + self.ff(mixed)
+        scores = self.attn_q(hidden).squeeze(-1).masked_fill(mask==0, float('-inf'))
+        return self.proj_out((hidden * torch.softmax(scores, dim=-1).unsqueeze(-1)).sum(1))
+
+class SeqFeatMLP(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.LayerNorm(32), nn.Linear(32, 128), nn.GELU(),
+            nn.Dropout(0.1), nn.Linear(128, 64), nn.LayerNorm(64)
+        )
+    def forward(self, x): return self.net(x)
+
+class PolyglotClassifierV3(nn.Module):
+    def __init__(self, base):
+        super().__init__()
+        self.encoder   = base
+        self.token_enc = PerTokenEncoder()
+        self.seq_mlp   = SeqFeatMLP()
+        fused = base.config.hidden_size + 256 + 64
+        self.classifier = nn.Sequential(
+            nn.LayerNorm(fused), nn.Linear(fused, 512),
+            nn.GELU(), nn.Dropout(0.2),
+            nn.Linear(512, 128), nn.GELU(),
+            nn.Dropout(0.1), nn.Linear(128, 1)
+        )
+    def forward(self, ids, mask, tf, sf):
+        hs  = self.encoder(input_ids=ids, attention_mask=mask).last_hidden_state
+        sem = (hs * mask.unsqueeze(-1)).sum(1) / mask.unsqueeze(-1).sum(1).clamp(min=1e-4)
+        return self.classifier(torch.cat([sem, self.token_enc(tf, mask), self.seq_mlp(sf)], dim=-1)).squeeze(-1)
+
+gen_tokenizer = AutoTokenizer.from_pretrained("Salesforce/codegen-350M-mono")
+if gen_tokenizer.pad_token is None:
+    gen_tokenizer.pad_token = gen_tokenizer.eos_token
+t5_tokenizer = RobertaTokenizer.from_pretrained("Salesforce/codet5-base", extra_special_tokens=None)
+
+gen_base      = AutoModelForCausalLM.from_pretrained("Salesforce/codegen-350M-mono", torch_dtype=torch.float16).to(DEVICE)
+metric_engine = PolyglotMetricEngine(gen_base).eval()
+
+t5_base  = T5EncoderModel.from_pretrained("Salesforce/codet5-base")
+detector = PolyglotClassifierV3(t5_base).to(DEVICE)
+
+weights_path = hf_hub_download(repo_id=REPO_ID, filename="model_weights.pt")
+detector.load_state_dict(torch.load(weights_path, map_location=DEVICE))
+detector.eval()
+
+def predict(code: str, threshold: float = THRESHOLD) -> dict:
+    if len(code.strip()) < 200:
+        return {"prediction": "Too Short", "ai_probability": None}
+
+    with torch.no_grad():
+        g      = gen_tokenizer(code, return_tensors="pt", padding="max_length", truncation=True, max_length=MAX_LEN).to(DEVICE)
+        tf, sf = metric_engine(g["input_ids"], g["attention_mask"])
+        t      = t5_tokenizer(code, return_tensors="pt", padding="max_length", truncation=True, max_length=MAX_LEN).to(DEVICE)
+        logits = detector(t["input_ids"], t["attention_mask"], tf.float(), sf.float())
+        prob   = torch.sigmoid(logits).item()
+
+    return {
+        "prediction": "AI Generated" if prob >= threshold else "Human Written",
+        "ai_probability": round(prob * 100, 2)
+    }