import json
import torch
from safetensors.torch import load_file
from huggingface_hub import snapshot_download

# See ai-detection-demo repo for the LinearDropoutHead class definition.
from gemma_4_e4b.head import LinearDropoutHead, score_from_logits

local = snapshot_download("DarrenJiaImbue/ai-detection-demo-gemma_4_e4b")
with open(f"{local}/config.json") as f:
    cfg = json.load(f)

head = LinearDropoutHead(
    vocab_size=cfg["vocab_size"], n_classes=cfg["n_classes"], dropout=cfg["dropout"]
)
head.load_state_dict(load_file(f"{local}/head.safetensors"))
head.eval()

# raw_logits: (batch, 262144) fp32 tensor of Gemma's last-token vocabulary logits.
with torch.no_grad():
    logits = head(raw_logits)
    bucket = logits.argmax(dim=-1)          # 0..3
    ai_score = score_from_logits(logits)    # continuous [0, 1]

Training data

Trained on DarrenJiaImbue/ai-detection-demo-gemma-logits — pre-computed int4 Gemma 4 E4B logits over DarrenJiaImbue/ai-detection-demo-dataset.

License

CC BY-NC-SA 4.0. Non-commercial research use only.