"""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()