File size: 8,535 Bytes
268e727 | 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 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | """Provider-independent voice state and synchronous finite-decision inference.
Hosts own async execution, cancellation, and device admission. Importing this
module does not import Studio, VoiceMem, Torch, Transformers, or ONNX Runtime.
"""
from dataclasses import dataclass
import json
import math
from pathlib import Path
import threading
import numpy as np
LABELS = ('ordinary', 'deep')
FIELDS = (('eot_probability', 1.), ('silence_s', 2.), ('span_s', 30.))
def observation_contract():
return {'version':1,'audio_fields':[name for name,_ in FIELDS],
'scales':[scale for _,scale in FIELDS],
'layout':['eot','silence','span','eot_present','silence_present','span_present','history_present'],
'text_preprocessing':'voicemem-decision-text',
'missing_value':0,'present_zero_has_mask':True}
def prepare(text, history=None, observations=None):
"""Normalize shared state; missing observations retain explicit masks."""
current = ' '.join(str(text).split())
if len(current) > 256:
current = current[:128] + ' … ' + current[-128:]
recent = []
for row in list(history or [])[-4:]:
role = row.get('role')
if role not in ('user', 'assistant'): continue
content = ' '.join(str(row.get('content') or '').split())[:80]
if content: recent.append(f'{role}: {content}')
state = {} if observations is None else observations
if not isinstance(state, dict) or set(state) - {name for name, _ in FIELDS}:
raise ValueError('Only shared physical observations and EOT probability are accepted')
values, masks = [], []
for name, scale in FIELDS:
value = state.get(name)
if value is None:
values.append(0.); masks.append(0.); continue
if (isinstance(value, bool) or not isinstance(value, (int, float))
or not math.isfinite(value) or value < 0
or (name == 'eot_probability' and value > 1)):
raise ValueError('Invalid shared observation')
values.append(min(value / scale, 1.)); masks.append(1.)
history_text = '\n'.join(recent)
return current, history_text, np.asarray(values + masks + [bool(history_text)], np.float32)
def collate(tokenizer, rows, *, tensors='pt'):
state = [prepare(row['text'], row.get('history'), row.get('audio_state')) for row in rows]
encoded = tokenizer([text for a,b,_ in state for text in (a,b)], padding=True,
truncation=True, max_length=384, return_tensors=tensors)
observed = np.stack([r[2] for r in state])
if tensors == 'pt':
import torch
observed = torch.from_numpy(observed)
return {'input_ids':encoded['input_ids'], 'attention_mask':encoded['attention_mask'], 'observed':observed}
def build_model(config, pretrained=None):
import torch
from torch import nn
from transformers import AutoModel
class StateDecisionModel(nn.Module):
def __init__(self):
super().__init__()
if config.model_type not in ('bert','modernbert'):
raise ValueError('Unsupported portable encoder')
options = {'trust_remote_code':False,
'attn_implementation':'sdpa' if config.model_type == 'modernbert' else 'eager'}
if config.model_type == 'modernbert': config.reference_compile = False
else: options['add_pooling_layer'] = False
self.encoder = (AutoModel.from_pretrained(str(pretrained),config=config,local_files_only=True,
weights_only=True,**options) if pretrained else AutoModel.from_config(config,**options))
self.hidden_size = config.hidden_size
width = 3*self.hidden_size + 7
self.classifier = nn.Sequential(nn.LayerNorm(width),nn.Linear(width,256),nn.GELU(),
nn.Dropout(.15),nn.Linear(256,2))
def forward(self,input_ids,attention_mask,observed):
hidden = self.encoder(input_ids=input_ids,attention_mask=attention_mask).last_hidden_state[:,0]
pair = hidden.reshape(observed.shape[0],2,self.hidden_size)
current = pair[:,0]
recent = pair[:,1] * observed[:,-1:]
return self.classifier(torch.cat((current,recent,current*recent,observed),dim=-1))
return StateDecisionModel()
def load_model(path, device='cpu', *, pretrained=False):
from transformers import AutoConfig, AutoTokenizer
from safetensors.torch import load_file
path = Path(path)
config_path = path if pretrained else path/'encoder'
config = AutoConfig.from_pretrained(config_path,local_files_only=True,trust_remote_code=False)
model = build_model(config,path if pretrained else None)
if not pretrained: model.load_state_dict(load_file(str(path/'model.safetensors')),strict=True)
tokenizer = AutoTokenizer.from_pretrained(path if pretrained else path/'tokenizer',
local_files_only=True,trust_remote_code=False,fix_mistral_regex=False)
return model.to(device).eval(),tokenizer
@dataclass(frozen=True)
class Decision:
label: str
p_ordinary: float
p_deep: float
calibration_scope: str = 'machine-supervised, not independently human-validated'
class Router:
"""Run a local model independently of ASR, TTS, LLM, and host application."""
def __init__(self,path,*,backend='onnx',device='cpu',admission=None):
if backend not in ('onnx','eager') or (backend=='onnx' and device!='cpu'):
raise ValueError('Unsupported runtime')
self.path = Path(path)
self.meta = json.loads((self.path/'depth_encoder.json').read_text())
if (self.meta.get('version')!=1 or self.meta.get('features_version')!=1
or self.meta.get('observation_contract')!=observation_contract()
or self.meta.get('task')!='voice-depth-encoder' or self.meta.get('labels')!=list(LABELS)
or self.meta.get('max_length')!=384 or self.meta.get('threshold')!=.5
or self.meta.get('history_policy')!='first80'):
raise ValueError('Incompatible portable decision contract')
self.temperature = float(self.meta['temperature'])
if not math.isfinite(self.temperature) or self.temperature<=0: raise ValueError('Invalid calibration')
self.backend,self.device,self.admission = backend,device,admission
self.model = self.tokenizer = None
self.lock = threading.Lock()
def _load(self):
if self.backend=='eager':
self.model,self.tokenizer = load_model(self.path,self.device)
else:
import onnxruntime as ort
from transformers import AutoTokenizer
self.tokenizer = AutoTokenizer.from_pretrained(self.path/'tokenizer',local_files_only=True,
trust_remote_code=False,fix_mistral_regex=False)
options = ort.SessionOptions()
options.intra_op_num_threads,options.inter_op_num_threads = 2,1
self.model = ort.InferenceSession(str(self.path/'model.onnx'),sess_options=options,
providers=['CPUExecutionProvider'])
def predict(self,text,*,history=None,observations=None):
from contextlib import nullcontext
current,recent,observed = prepare(text,history,observations)
if not current: return Decision('ordinary',1.,0.,'empty-input compatibility rule')
with self.lock, (self.admission() if self.admission else nullcontext()):
if self.model is None: self._load()
encoded = self.tokenizer([current,recent],padding=True,truncation=True,
max_length=384,return_tensors='np')
inputs = {'input_ids':encoded['input_ids'],'attention_mask':encoded['attention_mask'],
'observed':observed[None]}
if self.backend=='onnx': logits = self.model.run(['logits'],inputs)[0][0]
else:
import torch
with torch.inference_mode():
logits = self.model(**{k:torch.from_numpy(v).to(self.device) for k,v in inputs.items()})[0].float().cpu().numpy()
scaled = np.asarray(logits,np.float64)/self.temperature
if scaled.shape!=(2,) or not np.isfinite(scaled).all(): raise ValueError('Invalid decision logits')
probability = np.exp(scaled-scaled.max()); probability /= probability.sum()
deep = float(probability[1])
return Decision('deep' if deep>=.5 else 'ordinary',1.-deep,deep)
|