"""Score inert skill text with SkillScreen. Default decision threshold is 0.5."""
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, help='SKILL.md, or package JSON for the Package model')
    parser.add_argument('--model', default=str(Path(__file__).resolve().parent))
    parser.add_argument('--revision', default='main')
    args = parser.parse_args()
    if not torch.cuda.is_available():
        raise RuntimeError('This verified recipe requires CUDA with BF16 support')
    local = Path(args.model).is_dir()
    path = Path(args.model) / 'scoring_config.json' if local else Path(
        hf_hub_download(args.model, 'scoring_config.json', revision=args.revision))
    settings = json.loads(path.read_text())
    options = {'local_files_only': True} if local else {'revision':args.revision}
    tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=False, **options)
    assert tokenizer.num_special_tokens_to_add(False) == 0
    if settings['input']['scope'] == 'SKILL.md + related file contents':
        from package_inputs import encode_package
        row = json.loads(args.skill.read_text(encoding='utf-8'))
        ids, original, _ = encode_package(tokenizer, row,
            settings['input']['hypothesis'], settings['input']['max_length'])
    else:
        text = args.skill.read_text(encoding='utf-8')
        if len(text) > 262144:
            raise ValueError('For inputs above 262,144 characters, use the research pipeline bounded encoder')
        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 = len(ids)
        limit = settings['input']['max_length']
        if original > limit:
            ids = ids[:limit - len(suffix_ids)] + suffix_ids
    model = AutoModelForSequenceClassification.from_pretrained(args.model,
        dtype=torch.bfloat16, attn_implementation='sdpa', trust_remote_code=False, **options)
    assert model.config.label2id == {'contradiction':0,'entailment':1,'neutral':2}
    tokenizer.pad_token_id = model.config.pad_token_id
    model.config.use_cache = model.config.text_config.use_cache = False
    model = model.cuda().eval()
    inputs = torch.tensor([ids], dtype=torch.long, 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].float().double()
    if not torch.isfinite(logits).all():
        raise ValueError('Nonfinite model score')
    probability = logits.softmax(-1)[1].item()
    print(json.dumps({'malicious_probability':probability,'suspicious':probability >= 0.5,
        'truncated':original > settings['input']['max_length']}, indent=2))


if __name__ == '__main__':
    main()