113 lines
3.4 KiB
Python
113 lines
3.4 KiB
Python
import asyncio
|
|
import logging
|
|
import os
|
|
import re
|
|
from pathlib import Path
|
|
|
|
from faster_whisper import WhisperModel
|
|
from aiogram import Bot
|
|
|
|
from config import PROXY_ENABLED, PROXY_URL
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_model = None
|
|
|
|
BASE_DIR = Path(__file__).resolve().parent.parent.parent
|
|
MODELS_DIR = BASE_DIR / "models" / "faster-whisper"
|
|
VOICE_DIR = BASE_DIR / "bot" / "data" / "voice"
|
|
|
|
|
|
def _setup_proxy():
|
|
if PROXY_ENABLED and PROXY_URL:
|
|
os.environ["ALL_PROXY"] = PROXY_URL
|
|
logger.info("Proxy set for model download: %s", PROXY_URL[:30] + "...")
|
|
|
|
|
|
def _get_model() -> WhisperModel:
|
|
global _model
|
|
if _model is None:
|
|
MODELS_DIR.mkdir(parents=True, exist_ok=True)
|
|
VOICE_DIR.mkdir(parents=True, exist_ok=True)
|
|
model_path = str(MODELS_DIR / "base")
|
|
if not os.path.isdir(model_path) or not os.path.isfile(os.path.join(model_path, "model.bin")):
|
|
_setup_proxy()
|
|
logger.info("Model not found locally, downloading to %s...", MODELS_DIR)
|
|
_model = WhisperModel(
|
|
"base",
|
|
device="cpu",
|
|
cpu_threads=4,
|
|
compute_type="int8",
|
|
download_root=str(MODELS_DIR),
|
|
)
|
|
else:
|
|
logger.info("Loading model from %s...", model_path)
|
|
_model = WhisperModel(
|
|
model_path,
|
|
device="cpu",
|
|
cpu_threads=4,
|
|
compute_type="int8",
|
|
)
|
|
logger.info("Whisper model loaded")
|
|
return _model
|
|
|
|
|
|
async def download_voice(bot: Bot, file_id: str) -> str:
|
|
VOICE_DIR.mkdir(parents=True, exist_ok=True)
|
|
file = await bot.get_file(file_id)
|
|
path = str(VOICE_DIR / f"{file_id}.ogg")
|
|
await bot.download_file(file.file_path, destination=path)
|
|
logger.info("Voice downloaded: %s -> %s", file_id, path)
|
|
return path
|
|
|
|
|
|
async def convert_to_wav(input_path: str) -> str:
|
|
output_path = input_path.replace(".ogg", ".wav")
|
|
proc = await asyncio.create_subprocess_exec(
|
|
"ffmpeg", "-i", input_path, "-ar", "16000", "-ac", "1",
|
|
"-c:a", "pcm_s16le", output_path, "-y",
|
|
stdout=asyncio.subprocess.PIPE,
|
|
stderr=asyncio.subprocess.PIPE,
|
|
)
|
|
_, stderr = await proc.communicate()
|
|
if proc.returncode != 0:
|
|
logger.error(f"ffmpeg error: {stderr.decode(errors='replace')}")
|
|
raise RuntimeError("ffmpeg conversion failed")
|
|
logger.info(f"Converted to WAV: {output_path}")
|
|
return output_path
|
|
|
|
|
|
async def transcribe_audio(file_path: str) -> str:
|
|
model = _get_model()
|
|
loop = asyncio.get_running_loop()
|
|
|
|
def _transcribe():
|
|
segments, info = model.transcribe(file_path, language="ru", beam_size=5)
|
|
text = " ".join(seg.text for seg in segments)
|
|
return text.strip()
|
|
|
|
text = await loop.run_in_executor(None, _transcribe)
|
|
logger.info(f"Transcription result ({len(text)} chars): {text[:100]}...")
|
|
return text
|
|
|
|
|
|
def normalize_text(text: str) -> str:
|
|
text = text.strip()
|
|
if not text:
|
|
return text
|
|
text = text[0].upper() + text[1:]
|
|
if text[-1] not in ".!?…":
|
|
text += "."
|
|
text = re.sub(r"\s+", " ", text)
|
|
return text
|
|
|
|
|
|
async def cleanup_files(*paths: str):
|
|
for path in paths:
|
|
if path and os.path.exists(path):
|
|
try:
|
|
os.unlink(path)
|
|
logger.debug(f"Cleaned up: {path}")
|
|
except Exception as e:
|
|
logger.warning(f"Cleanup failed for {path}: {e}")
|