File size: 6,087 Bytes
4fa3363 | 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 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """Local inference with versioned SkillScreen deployment models."""
import json
import math
from pathlib import Path
from .package_inputs import encode_package
ROOT = Path(__file__).resolve().parents[1]
HYPOTHESIS = 'This AI agent skill is malicious.'
def encode_input(tokenizer, row, mode, max_length):
if mode not in ('package', 'markdown'):
raise ValueError('Choose package or markdown input mode.')
if tokenizer.num_special_tokens_to_add(False) != 0:
raise ValueError('The released model requires a tokenizer without added special tokens.')
if mode == 'package':
return encode_package(tokenizer, row, HYPOTHESIS, max_length)
suffix = '\nHypothesis: ' + HYPOTHESIS
suffix_ids = tokenizer.encode(suffix, add_special_tokens=False)
if type(max_length) is not int or max_length <= len(suffix_ids) + 8:
raise ValueError('Input budget is too small for the premise and fixed hypothesis.')
ids = tokenizer.encode('Premise: ' + row['skill_text'] + suffix, add_special_tokens=False)
original = len(ids)
if original > max_length:
ids = ids[:max_length - len(suffix_ids)] + suffix_ids
return ids, original, {'original_tokens': original, 'used_tokens': len(ids),
'truncated': original > max_length, 'files': ['SKILL.md']}
def load_release(mode, model_dir=None):
"""Resolve one local release or the pinned Hub release for this input mode."""
if model_dir is None and (ROOT / 'deployment_settings.json').is_file():
model_dir = ROOT
if model_dir is not None:
source = str(Path(model_dir).resolve())
settings = json.loads((Path(source) / 'deployment_settings.json').read_text(encoding='utf-8'))
options = {'local_files_only': True, 'trust_remote_code': False}
revision = settings['release']
else:
from huggingface_hub import hf_hub_download
spec = json.loads((ROOT / 'configs/community-models.json').read_text(encoding='utf-8'))[mode]
source, revision = spec['repo_id'], spec['revision']
path = hf_hub_download(source, 'deployment_settings.json', revision=revision)
settings = json.loads(Path(path).read_text(encoding='utf-8'))
if settings['repo_id'] != source:
raise ValueError('Model identity differs from the pinned release.')
options = {'revision': revision, 'trust_remote_code': False}
if settings['input_mode'] != {'package': 'package32k', 'markdown': 'markdown8k'}[mode]:
raise ValueError('This model uses a different input mode.')
if settings['hypothesis'] != HYPOTHESIS:
raise ValueError('Unexpected hypothesis in deployment settings.')
return source, options, settings, revision
def select_threshold(settings, operating_point='auto'):
policy = settings['default_policy'] if operating_point == 'auto' else operating_point
key = {'native': 'default', 'high-recall': 'high_recall', 'balanced': 'balanced'}.get(policy)
if key is None:
raise ValueError('Unknown operating point.')
gate = settings['thresholds'][key]
raw, native = gate['raw_threshold'], gate['native_threshold']
if not all(math.isfinite(x) for x in (raw, native)) or not 0 < native < 1:
raise ValueError('Invalid saved threshold.')
if not math.isclose(raw, math.log(native) - math.log1p(-native), abs_tol=1e-9):
raise ValueError('Raw and native thresholds disagree.')
return policy, gate
def score_model(row, mode, model_dir=None, operating_point='auto'):
try:
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
except ImportError as exc:
raise ValueError('Model inference needs requirements/inference.txt. Package inspection and dependency checks need only Python.') from exc
if not torch.cuda.is_available() or not torch.cuda.is_bf16_supported():
raise ValueError('The released inference recipe requires a CUDA GPU with BF16 support. Use --inspect or --checks dependencies on this machine.')
source, options, settings, revision = load_release(mode, model_dir)
policy, gate = select_threshold(settings, operating_point)
torch.set_num_threads(8)
torch.manual_seed(42)
tokenizer = AutoTokenizer.from_pretrained(source, **options)
ids, _, coverage = encode_input(tokenizer, row, mode, settings['max_length'])
model = AutoModelForSequenceClassification.from_pretrained(
source, dtype=torch.bfloat16, attn_implementation='sdpa', **options)
if model.config.label2id != {'contradiction': 0, 'entailment': 1, 'neutral': 2}:
raise ValueError('Unexpected model label mapping.')
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('Model returned nonfinite logits.')
score = logits.softmax(-1)[1].item()
raw = (logits[1] - torch.logsumexp(logits[[0, 2]], dim=0)).item()
return {'status': 'complete', 'repo_id': settings['repo_id'], 'revision': revision,
'release': settings['release'], 'training_scope': settings['training_scope'],
'dataset': settings['dataset']['repo_id'], 'dataset_revision': settings['dataset']['revision'],
'protocol': 'three_fold_oof_then_all_data_refit', 'operating_point': policy,
'threshold_source': 'native 0.5 reference' if policy == 'native' else 'frozen three-fold OOF threshold',
'input_mode': mode, 'score': score, 'score_kind': 'native entailment softmax',
'raw_log_odds': raw, 'threshold': gate['native_threshold'],
'raw_threshold': gate['raw_threshold'], 'review': raw >= gate['raw_threshold'],
'coverage': coverage}
|