| """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} |
|
|
