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