| """Score one SKILL.md with the full SkillJev-2B model; never execute its text.""" |
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import torch |
| from huggingface_hub import hf_hub_download |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("skill", type=Path) |
| parser.add_argument("--model", default="SkillJev/SkillJev-2B") |
| parser.add_argument("--revision", default="v1") |
| args = parser.parse_args() |
| if not torch.cuda.is_available(): |
| raise RuntimeError("This release recipe requires CUDA with BF16 support") |
| total = torch.cuda.get_device_properties(0).total_memory |
| torch.cuda.set_per_process_memory_fraction(min(1., 8 * 1024**3 / total)) |
| local = Path(args.model).is_dir() |
| config = Path(args.model) / "scoring_config.json" if local else Path( |
| hf_hub_download(args.model, "scoring_config.json", revision=args.revision)) |
| settings = json.loads(config.read_text()) |
| options = {"local_files_only": True} if local else {"revision": args.revision} |
| tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=False, **options) |
| model = AutoModelForSequenceClassification.from_pretrained(args.model, |
| dtype=torch.bfloat16, attn_implementation="sdpa", trust_remote_code=False, **options) |
| assert model.config.label2id == settings["label2id"] |
| assert tokenizer.num_special_tokens_to_add(False) == 0 |
| tokenizer.pad_token_id = model.config.pad_token_id |
| model.config.use_cache = model.config.text_config.use_cache = False |
| model = model.cuda().eval() |
| text = args.skill.read_text(encoding="utf-8") |
| suffix = "\nHypothesis: " + settings["input"]["hypothesis"] |
| suffix_ids = tokenizer.encode(suffix, add_special_tokens=False) |
| ids = tokenizer.encode("Premise: " + text + suffix, add_special_tokens=False) |
| original_tokens = len(ids) |
| limit = settings["input"]["max_length"] |
| if len(ids) > limit: |
| ids = ids[:limit - len(suffix_ids)] + suffix_ids |
| inputs = torch.tensor([ids], device="cuda") |
| with torch.inference_mode(), torch.autocast("cuda", dtype=torch.bfloat16): |
| logits = model(input_ids=inputs, attention_mask=torch.ones_like(inputs)).logits[0] |
| logits = logits.float().cpu().double() |
| assert torch.isfinite(logits).all() |
| z = logits[1] - torch.logsumexp(logits[[0, 2]], dim=0) |
| calibration = settings["calibration"] |
| p = torch.sigmoid(calibration["slope"] * z + calibration["intercept"]).item() |
| decision = settings["decision"] |
| print(json.dumps({"skill_id": args.skill.name, "malicious_probability": p, |
| "threshold": decision["calibrated_probability_threshold"], |
| "review": z.item() >= decision["raw_log_odds_threshold"], |
| "truncated": original_tokens > limit}, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
