File size: 2,455 Bytes
a1572a3 d1840ef a1572a3 d1840ef a1572a3 d1840ef a1572a3 d1840ef | 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 | # SPDX-License-Identifier: CC-BY-NC-SA-4.0
# Copyright 2026 TeraSpace.
# Edit these variables, then run this file. No command-line arguments.
MODEL_ID = 'TeraSpace/TeraAnaly' # Or a local folder containing config + weights.
AUDIO_PATH = 'audio.wav' # Or ['first.wav', 'second.wav', ...].
TEXT = 'Конечно, помню! Мы так здорово проводили время!' # Or matching texts.
DEVICE = 'auto'
BATCH_SIZE = 4
import json
from pathlib import Path
import torch
from modeling_tera_analy import TeraAnaly, read_wav
if __name__ == '__main__':
torch.set_num_threads(4)
batched = isinstance(AUDIO_PATH, (list, tuple))
paths = list(AUDIO_PATH) if batched else [AUDIO_PATH]
if batched and (not isinstance(TEXT, (list, tuple)) or len(TEXT) != len(paths)):
raise ValueError('Supply one TEXT per AUDIO_PATH')
texts = list(TEXT) if batched else [TEXT]
if any(not isinstance(text, str) or not text.strip() for text in texts):
raise ValueError('TEXT must contain the matching transcript for each recording')
if not isinstance(BATCH_SIZE, int) or BATCH_SIZE < 1:
raise ValueError('BATCH_SIZE must be a positive integer')
device = ('cuda' if torch.cuda.is_available() else 'cpu') if DEVICE == 'auto' else DEVICE
if Path(MODEL_ID).is_dir():
model = TeraAnaly.from_pretrained(MODEL_ID).to(device)
else:
from transformers import AutoModel
model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True).to(device)
results = []
for start in range(0, len(paths), BATCH_SIZE):
# Decode each recording once; one frozen encoder pass per batch.
waveforms = [read_wav(Path(__file__).resolve().parent / path) for path in paths[start:start+BATCH_SIZE]]
states = model.encoder(waveforms)
# Each head reuses states. Stress runs our batched CTC internally.
outputs = {
'age': model.age_head(states),
'gender': model.gender_head(states),
'emotion': model.emotion_head(states),
'stressed_text': model.stress_head(states, texts[start:start+BATCH_SIZE]),
}
results.extend(dict(zip(outputs, values)) for values in zip(*outputs.values()))
# Optional character logits from the same states: model.ctc_head(states)
del states, waveforms
print(json.dumps(results if batched else results[0], ensure_ascii=False, indent=2, allow_nan=False))
|