File size: 3,152 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 | """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
|