66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
"""Движки распознавания речи: whisper-large-v3-russian и Qwen3-ASR-1.7B."""
|
|
import warnings
|
|
import logging
|
|
|
|
import librosa
|
|
import numpy as np
|
|
import soundfile as sf
|
|
import torch
|
|
import torchaudio
|
|
from transformers import WhisperForConditionalGeneration, WhisperProcessor
|
|
|
|
from qwen_asr import Qwen3ASRModel
|
|
|
|
from .config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class WhisperModel:
|
|
"""whisper-large-v3-russian (transformers)."""
|
|
|
|
def __init__(self, device):
|
|
self.model = WhisperForConditionalGeneration.from_pretrained(
|
|
settings.WHISPER_MODEL_PATH
|
|
)
|
|
self.processor = WhisperProcessor.from_pretrained(settings.WHISPER_MODEL_PATH)
|
|
self.device = device
|
|
self.model.to(device)
|
|
|
|
def transcribe(self, audio_path):
|
|
wav, sr = sf.read(audio_path)
|
|
wav = wav.mean(axis=1) if wav.ndim > 1 else wav
|
|
wav16 = torchaudio.functional.resample(
|
|
torch.from_numpy(np.asarray(wav, dtype=np.float32)), sr, 16000
|
|
)
|
|
inputs = self.processor(
|
|
wav16, sampling_rate=16000, return_tensors="pt"
|
|
)
|
|
with torch.no_grad():
|
|
ids = self.model.generate(inputs["input_features"].to(self.device))
|
|
return self.processor.batch_decode(ids.cpu().numpy(), skip_special_tokens=True)
|
|
|
|
|
|
class QwenModel:
|
|
"""Qwen3-ASR-1.7B (qwen-asr)."""
|
|
|
|
def __init__(self, device):
|
|
self.model = Qwen3ASRModel.from_pretrained(settings.QWEN_MODEL_NAME)
|
|
self.device = device
|
|
|
|
def transcribe(self, audio_path):
|
|
result = self.model.transcribe(audio_path)
|
|
return result.get("text") if isinstance(result, dict) else result
|
|
|
|
|
|
def load_models(device):
|
|
"""Загрузить все активные модели."""
|
|
models = {}
|
|
if settings.LOAD_WHISPER:
|
|
logger.info("Загрузка whisper-large-v3-russian...")
|
|
models["whisper"] = WhisperModel(device)
|
|
if settings.LOAD_QWEN:
|
|
logger.info("Загрузка Qwen3-ASR-1.7B...")
|
|
models["qwen3"] = QwenModel(device)
|
|
return models
|