model.py

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()) < 120:
        return {"pred": "too short", "prob": 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 {
        "pred": "ai-generated" if prob >= threshold else "human",
        "prob": (prob * 100, 2)
    }