"""Single-pass OpenVINO typed decisions using the frozen LFM12 prompt and readout.""" from __future__ import annotations import json from pathlib import Path import time import numpy as np from transformers import AutoTokenizer from data_io import encode_example, label_ids, check_prompt_format def artifact_prompt_format(graph): directory = Path(graph).parent formats = [] for name in ('compression.json', 'export.json'): if (directory / name).is_file(): formats.append(check_prompt_format(json.loads((directory / name).read_text()).get('prompt_format', 'canonical'))) if len(set(formats)) > 1: raise ValueError('Artifact prompt format metadata conflicts') return formats[0] if formats else 'canonical' def typed_example(state, question): kind = question['type'] criteria = question.get('criteria') if kind == 'noul': labels, options = ['no', 'yes'], ['No', 'Yes'] if criteria is not None: if not isinstance(criteria, dict): raise ValueError('Noul criteria must map false and true') criteria = {str(k).lower(): v for k, v in criteria.items()} if set(criteria) != {'false', 'true'}: raise ValueError('Noul criteria must map exactly false and true') options = [label + ': ' + (criteria[key] if isinstance(criteria[key], str) else json.dumps(criteria[key], ensure_ascii=False)) for label, key in [('No', 'false'), ('Yes', 'true')]] elif kind == 'choice': if isinstance(criteria, dict): labels = [str(k) for k in criteria] options = [str(k) + (': ' + (v if isinstance(v, str) else json.dumps(v, ensure_ascii=False)) if v is not None else '') for k, v in criteria.items()] elif isinstance(criteria, list): labels = [str(k) for k in criteria] options = list(labels) else: raise ValueError('Choice requires criteria') elif kind == 'score': if not isinstance(criteria, list): raise ValueError('Score requires ordered criteria') labels = [str(i) for i in range(len(criteria))] options = [f'{i}: ' + (v if isinstance(v, str) else json.dumps(v, ensure_ascii=False)) for i, v in enumerate(criteria)] else: raise ValueError('Supported types are choice, score and noul') if len(set(labels)) != len(labels): raise ValueError('Labels must be distinct') return {'state': state, 'instructions': question['instructions'], 'options': options}, labels class Runtime: def __init__(self, graph, tokenizer, threads=16, precision='bf16', buckets=(128,256,512), short_threads=None, short_length=128): import openvino as ov self.graph = Path(graph) self.prompt_format = artifact_prompt_format(graph) if (short_threads is not None and (type(short_threads) is not int or short_threads < 1)) or type(short_length) is not int or short_length < 1: raise ValueError('Short thread count/length must be positive integers') self.short_threads, self.short_length = short_threads, short_length self.tokenizer = AutoTokenizer.from_pretrained(str(tokenizer), local_files_only=True) self.codes = np.array(label_ids(self.tokenizer), dtype=np.int64) self.buckets = tuple(buckets) self.compiled = {} self.requests = {} self.core = ov.Core() self.base_graph = self.core.read_model(str(self.graph)) self.config = {'PERFORMANCE_HINT':'LATENCY','INFERENCE_PRECISION_HINT':precision, 'INFERENCE_NUM_THREADS':threads,'NUM_STREAMS':1,'ENABLE_CPU_PINNING':True, 'ENABLE_HYPER_THREADING':False} def compile(self, length): if length not in self.compiled: model = self.base_graph.clone() names = {p.get_any_name() for p in model.inputs} shapes = {'input_ids':[1,length],'attention_mask':[1,length]} if 'selected_token_ids' in names: shapes['selected_token_ids'] = [26] model.reshape(shapes) config = dict(self.config) if self.short_threads is not None and length <= self.short_length: config['INFERENCE_NUM_THREADS'] = self.short_threads compiled = self.core.compile_model(model,'CPU',config) self.compiled[length] = compiled self.requests[length] = compiled.create_infer_request() return self.requests[length] def arrays(self, row, length): ids = np.full((1,length), self.tokenizer.pad_token_id,dtype=np.int64) mask = np.zeros_like(ids) ids[0,:len(row['input_ids'])] = row['input_ids'] mask[0,:len(row['input_ids'])] = 1 values = {'input_ids':ids, 'attention_mask':mask, 'selected_token_ids':self.codes} return {p.get_any_name(): values[p.get_any_name()] for p in self.compiled[length].inputs} def warmup(self): row = encode_example(self.tokenizer, {'state':'A red box is on the table.', 'instructions':'Which color is the box?', 'options':['Red','Blue']},prompt_format=self.prompt_format) timings = {} for length in self.buckets: start = time.perf_counter() request = self.compile(length) for _ in range(3): request.infer(self.arrays(row,length)) timings[length] = time.perf_counter()-start return timings def predict(self, state, questions): answers, audits = {}, {} for qid, question in questions.items(): ex, labels = typed_example(state,question) row = encode_example(self.tokenizer,ex,max_len=max(self.buckets),prompt_format=self.prompt_format) length = next(b for b in self.buckets if len(row['input_ids']) <= b) request = self.compile(length) request.infer(self.arrays(row,length)) logits = request.get_output_tensor(0).data.copy()[0,:len(labels)].astype(np.float64) if not np.isfinite(logits).all(): raise RuntimeError('Nonfinite decision logits') probabilities = np.exp(logits-logits.max()) probabilities /= probabilities.sum() mapping = dict(zip(labels, probabilities.tolist())) kind = question['type'] answer = {'type':kind,'probabilities':mapping} if kind == 'noul': answer['noul'] = mapping['yes'] elif kind == 'choice': answer['choice'] = labels[int(probabilities.argmax())] else: answer['score'] = sum(i*p for i,p in enumerate(probabilities.tolist())) answers[qid], audits[qid] = answer, row['audit'] return {'answers':answers,'input_audit':audits,'probabilities_calibrated':False}