File size: 3,131 Bytes
b646fbd | 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 52 53 54 55 56 57 58 59 60 | """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()
|