gigaAM / app.py
recycleactor's picture
Upload 2 files
2072e4a verified
Raw
History Blame Contribute Delete
8.77 kB
import glob
import os
import subprocess
import tempfile
# ВАЖНО: для Space на тарифе ZeroGPU HF теперь при старте контейнера
# сканирует код на наличие хотя бы одной функции с @spaces.GPU — если её
# нет, падает "Runtime error: No @spaces.GPU function detected during
# startup" (раньше хватало просто импорта spaces, сейчас уже нет).
# Реальная модель у нас работает на CPU, GPU не нужен, поэтому заводим
# функцию-пустышку только ради прохождения этой проверки — она нигде не
# вызывается, так что GPU-квота ZeroGPU не тратится.
import spaces
@spaces.GPU
def _zerogpu_startup_probe():
"""Никогда не вызывается — нужна только чтобы Space прошёл валидацию
ZeroGPU на старте (см. комментарий выше)."""
return None
import gigaam
import gradio as gr
from faster_whisper import WhisperModel
from faster_whisper.audio import decode_audio
# Модель грузим ОДИН раз при старте контейнера, не на каждый запрос.
# v3_e2e_ctc — версия с пунктуацией и нормализацией текста (не голый поток
# слов без знаков препинания) и CTC-декодер (быстрее RNNT, для транскрибации
# голосовых в чате этого достаточно).
# Никакого @spaces.GPU нигде не используем — Space создан на тарифе ZeroGPU
# (сейчас у бесплатных аккаунтов HF нет варианта CPU Basic при создании),
# но раз GPU нигде не запрашивается, приложение работает на обычном CPU
# этого же контейнера и не трогает дневную квоту GPU-секунд. GigaAM
# компактная (220-240М параметров), GPU ей и не нужен.
MODEL = gigaam.load_model("v3_e2e_ctc", device="cpu")
# GigaAM обучена только на русской речи — на других языках она не "не
# понимает" в смысле пустого результата, а галлюцинирует русский текст по
# созвучию (особенно опасно для карточки "о чём песня" — см. index.js:
# ошибочная "русская" расшифровка иностранного текста песни ломает весь
# последующий ответ модели). Поэтому перед прогоном через GigaAM отдельно
# определяем язык лёгкой моделью faster-whisper (tiny, только
# language-id — полную транскрибацию ею не делаем, для этого дальше в
# пайплайне бота уже есть Groq Whisper). tiny достаточно для language-id:
# точность определения языка у tiny почти не хуже, чем у крупных моделей,
# в отличие от точности самой транскрибации.
LANG_ID_MODEL = WhisperModel("tiny", device="cpu", compute_type="int8")
# Ниже этого порога уверенности определённому языку не доверяем и всё
# равно отдаём на GigaAM — так безопаснее для коротких/тихих голосовых,
# где language-id менее надёжен, чем ошибочно отбраковывать русскую речь.
LANG_ID_MIN_CONFIDENCE = 0.5
def _detect_language(wav_path: str) -> tuple[str, float]:
"""Определяет язык по первому 30-секундному сегменту WAV (больше и не
нужно — language-id у Whisper всегда смотрит только на первый сегмент,
длинное аудио тут не даёт точности, только замедляет)."""
audio = decode_audio(wav_path, sampling_rate=16000)
language, probability, _ = LANG_ID_MODEL.detect_language(audio)
return language, probability
# Нативный model.transcribe() ограничен 25 секундами аудио за вызов.
# Официальный обход (.transcribe_longform) тянет за собой pyannote.audio
# и токен Hugging Face — ради простой последовательной нарезки заводить
# это не стоит, поэтому режем длинные голосовые сами через ffmpeg.
CHUNK_SECONDS = 20
def _to_wav(src_path: str) -> str:
"""Конвертирует присланный файл (у Telegram голосовые — .ogg/opus, но
подойдёт любой формат, который понимает ffmpeg — mp3, flac и т.д.) в
16kHz mono WAV. Модели нужен предсказуемый формат на входе."""
dst_path = tempfile.mktemp(suffix=".wav")
subprocess.run(
["ffmpeg", "-y", "-i", src_path, "-ar", "16000", "-ac", "1", dst_path],
check=True,
capture_output=True,
)
return dst_path
def _split_into_chunks(wav_path: str) -> list[str]:
"""Режет WAV на куски по CHUNK_SECONDS через ffmpeg segment muxer.
Простая последовательная нарезка без перехлёста — изредка может
срезать слово ровно на границе куска, но для голосовых в чате это
приемлемый компромисс против того, чтобы тянуть pyannote+HF-токен
только ради длинных сообщений."""
out_dir = tempfile.mkdtemp()
pattern = os.path.join(out_dir, "chunk_%03d.wav")
subprocess.run(
[
"ffmpeg", "-y", "-i", wav_path,
"-f", "segment", "-segment_time", str(CHUNK_SECONDS),
"-c", "copy", pattern,
],
check=True,
capture_output=True,
)
return sorted(glob.glob(os.path.join(out_dir, "chunk_*.wav")))
def transcribe(audio_path: str, force_giga: bool = False) -> str:
"""force_giga=True пропускает language-id и всегда гонит через GigaAM —
для ручного выбора движка на стороне бота ("Жень текст giga"). В
остальных случаях (auto) сначала проверяем язык и на не-русском сразу
возвращаем "" — вызывающий код в боте трактует пустой ответ как "GigaAM
не справился" и молча уходит на Groq Whisper (см. transcribeViaGigaAM в
index.js), поэтому отдельный сигнал "чужой язык" наружу прокидывать не
нужно, пустой строки достаточно."""
if not audio_path:
return ""
wav_path = _to_wav(audio_path)
chunks: list[str] = []
try:
if not force_giga:
language, probability = _detect_language(wav_path)
if language != "ru" or probability < LANG_ID_MIN_CONFIDENCE:
return ""
chunks = _split_into_chunks(wav_path)
parts = [str(MODEL.transcribe(chunk)).strip() for chunk in chunks]
return " ".join(p for p in parts if p).strip()
finally:
os.remove(wav_path)
for c in chunks:
try:
os.remove(c)
except OSError:
pass
if chunks:
try:
os.rmdir(os.path.dirname(chunks[0]))
except OSError:
pass
# api_name фиксирует имя эндпоинта — так со стороны бота адрес вызова
# предсказуем (/call/transcribe) и не завязан на автогенерируемое имя.
demo = gr.Interface(
fn=transcribe,
inputs=[
gr.Audio(type="filepath", label="Аудио"),
gr.Checkbox(value=False, label="Принудительно GigaAM (без проверки языка)"),
],
outputs=gr.Textbox(label="Транскрипция"),
title="GigaAM-v3 — распознавание речи (RU)",
api_name="transcribe",
)
if __name__ == "__main__":
demo.launch()