# 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))