File size: 1,828 Bytes
f3d5a27 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | import json
from pathlib import Path
import numpy as np
from peft import PeftModel
from safetensors.numpy import load_file
from sentence_transformers import SentenceTransformer
class MuncheAIDetector:
def __init__(self, model_dir: str | Path = "."):
model_dir = Path(model_dir)
self.encoder = SentenceTransformer("google/embeddinggemma-300m")
self.encoder[0].auto_model = PeftModel.from_pretrained(
self.encoder[0].auto_model,
model_dir,
)
self.encoder.max_seq_length = 2048
head = load_file(model_dir / "linear_head.safetensors")
self.weight = head["linear.weight"].reshape(-1)
self.bias = float(head["linear.bias"].item())
calibration = json.loads(
(model_dir / "calibration.json").read_text(encoding="utf-8")
)
self.human_max = calibration["thresholds"]["human_max"]
self.ai_min = calibration["thresholds"]["ai_min"]
def predict(self, text: str) -> dict:
tokens = self.encoder.tokenizer.encode(text, add_special_tokens=False)
if not 384 <= len(tokens) <= 2048:
raise ValueError("Input must contain 384 to 2,048 EmbeddingGemma tokens.")
embedding = self.encoder.encode(
[text],
normalize_embeddings=True,
convert_to_numpy=True,
)[0]
logit = float(embedding @ self.weight + self.bias)
score = (
float(1.0 / (1.0 + np.exp(-logit)))
if logit >= 0
else float(np.exp(logit) / (1.0 + np.exp(logit)))
)
label = (
"human"
if score <= self.human_max
else "llm"
if score >= self.ai_min
else "uncertain"
)
return {"label": label, "score": score, "tokens": len(tokens)}
|