| """Multilingual encoder with text, context and speech-state fusion.""" |
| from pathlib import Path |
|
|
| import torch |
| from torch import nn |
| from transformers import AutoConfig, AutoModel, AutoTokenizer |
| from safetensors.torch import load_file |
|
|
| from .features import collate |
|
|
|
|
| class VoiceEncoder(nn.Module): |
| def __init__(self, config, *, pretrained=None, pooling='cls', dropout=.15): |
| super().__init__() |
| if pooling not in ('cls', 'mean'): |
| raise ValueError('Unsupported pooling') |
| options = {'trust_remote_code': False, |
| 'attn_implementation': 'sdpa' if config.model_type == 'modernbert' else 'eager'} |
| if config.model_type == 'modernbert': |
| config.reference_compile = False |
| if config.model_type == 'bert': |
| options['add_pooling_layer'] = False |
| self.encoder = (AutoModel.from_pretrained(pretrained, config=config, local_files_only=True, |
| weights_only=True, **options) |
| if pretrained else AutoModel.from_config(config, **options)) |
| self.pooling, self.hidden_size = pooling, config.hidden_size |
| width = 3*self.hidden_size+7 |
| self.classifier = nn.Sequential(nn.LayerNorm(width), nn.Linear(width, 256), nn.GELU(), |
| nn.Dropout(dropout), 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 |
| if self.pooling == 'mean': |
| mask = attention_mask.unsqueeze(-1).to(hidden.dtype) |
| states = (hidden*mask).sum(dim=1)/mask.sum(dim=1).clamp(min=1) |
| else: |
| states = hidden[:, 0] |
| pairs = states.reshape(observed.shape[0], 2, self.hidden_size) |
| current = pairs[:, 0] |
| recent = pairs[:, 1]*observed[:, -1:] |
| return self.classifier(torch.cat((current, recent, current*recent, observed), dim=-1)) |
|
|
|
|
| def tokenizer_at(path): |
| return AutoTokenizer.from_pretrained(path, local_files_only=True, trust_remote_code=False, |
| fix_mistral_regex=False) |
|
|
|
|
| def load(path, metadata, device='cpu'): |
| path = Path(path) |
| config = AutoConfig.from_pretrained(path/'encoder', local_files_only=True, trust_remote_code=False) |
| model = VoiceEncoder(config, pooling=metadata['pooling']) |
| model.load_state_dict(load_file(str(path/'model.safetensors')), strict=True) |
| return model.to(device).eval(), tokenizer_at(path/'tokenizer') |
|
|
|
|
| def predict(model, tokenizer, rows, *, device='cpu', batch_size=16, max_length=384, drop_audio=False, |
| history_policy='first80'): |
| model.eval() |
| logits = [] |
| with torch.inference_mode(): |
| for start in range(0, len(rows), batch_size): |
| inputs = collate(tokenizer, rows[start:start+batch_size], max_length, drop_audio=drop_audio, |
| history_policy=history_policy) |
| values = model(**{key: value.to(device) for key, value in inputs.items()}) |
| logits.extend(values.float().cpu().tolist()) |
| return logits |
|
|
