ml_service/main.py - python-ml-service - Dmitry-F/cbt_tool

ml_service/main.py
134 строки5 KB
Dmitriy-Filatov
chore: add distortions.py and update ML service configs
14 апр 2026, 03:46
14 апр 2026, 03:46f5819c2
import asyncio
import logging
from io import BytesIO
from fastapi import FastAPI, UploadFile, File, HTTPException, Request
from faster_whisper import WhisperModel
from transformers import pipeline
from prometheus_client import Counter, Histogram, generate_latest
import time

from distortions import detect_distortions, format_distortions

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

REQUESTS = Counter('analyze_requests_total', 'Total requests')
ERRORS = Counter('analyze_errors_total', 'Error count')
DURATION = Histogram('analyze_duration_seconds', 'Processing time')

MAX_FILE_SIZE = 10 * 1024 * 1024
TIMEOUT_SECONDS = 30
CONFIDENCE_THRESHOLD = 0.6

EMOTION_MAP = {
    "neutral": "Нейтрально",
    "joy": "Радость",
    "sadness": "Грусть",
    "anger": "Гнев",
    "fear": "Страх / Тревога",
    "disgust": "Отвращение",
    "surprise": "Удивление",
    "guilt": "Вина",
    "shame": "Стыд",
    "interest": "Интерес"
}

def map_emotion_label(label: str) -> str:
    return EMOTION_MAP.get(label, label)

print("Loading Whisper base...")
whisper_model = WhisperModel("base", device="cpu", compute_type="int8")
print("Loading ruBERT (Aniemore 9 emotions)...")
classifier = pipeline("sentiment-analysis", model="Aniemore/rubert-tiny2-russian-emotion-detection")

def convert_to_wav(audio_bytes: bytes) -> bytes:
    import subprocess
    process = subprocess.Popen(
        ['ffmpeg', '-i', 'pipe:0', '-f', 'wav', '-acodec', 'pcm_s16le', '-ac', '1', '-ar', '16000', '-y', 'pipe:1'],
        stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )
    out, err = process.communicate(input=audio_bytes)
    if process.returncode != 0:
        logger.error(f"FFmpeg error: {err.decode()}")
        raise RuntimeError("FFmpeg conversion failed")
    return out

def process_audio(audio_bytes: bytes):
    try:
        wav_bytes = convert_to_wav(audio_bytes)
        audio_file = BytesIO(wav_bytes)
        segments, _ = whisper_model.transcribe(audio_file, language="ru")
        text = " ".join([seg.text for seg in segments])
        if not text.strip():
            return None, "no_speech"
        distortions_dict = detect_distortions(text)
        distortions_list = format_distortions(distortions_dict)
        analysis = classifier(text[:512])
        label = analysis[0]['label']
        confidence = analysis[0]['score']
        if confidence < CONFIDENCE_THRESHOLD:
            emotion_name = "Смешанные / Неопределённо"
        else:
            emotion_name = map_emotion_label(label)
        return {
            "transcription": text,
            "distortions": distortions_list,
            "emotion": {
                "primary": emotion_name,
                "secondary": None,
                "intensity": round(confidence * 100, 2),
                "confidence": round(confidence, 3)
            },
            "mismatch": {"detected": False, "description": None},
            "processing_time_ms": 0
        }, None
    except Exception as e:
        logger.error(f"Process audio error: {str(e)}")
        return None, str(e)

@app.post("/analyze")
async def analyze_audio(request: Request, file: UploadFile = File(...)):
    REQUESTS.inc()
    start = time.time()
    content_length = request.headers.get('content-length')
    if content_length and int(content_length) > MAX_FILE_SIZE:
        ERRORS.inc()
        raise HTTPException(413, {"error": "file_too_large"})
    try:
        audio_bytes = await file.read()
        if len(audio_bytes) < 2000:
            ERRORS.inc()
            raise HTTPException(422, {"error": "empty_audio"})
        if len(audio_bytes) > MAX_FILE_SIZE:
            ERRORS.inc()
            raise HTTPException(413, {"error": "file_too_large"})
        loop = asyncio.get_running_loop()
        result, error = await asyncio.wait_for(
            loop.run_in_executor(None, process_audio, audio_bytes),
            timeout=TIMEOUT_SECONDS
        )
        if error:
            ERRORS.inc()
            raise HTTPException(422, {"error": error})
        result["processing_time_ms"] = round((time.time() - start) * 1000)
        logger.info(f"Processed {len(audio_bytes)} bytes in {result['processing_time_ms']} ms")
        return result
    except asyncio.TimeoutError:
        ERRORS.inc()
        raise HTTPException(504, {"error": "timeout"})
    except HTTPException:
        raise
    except Exception as e:
        ERRORS.inc()
        logger.error(f"Unexpected error: {str(e)}")
        raise HTTPException(500, {"error": "internal", "message": str(e)})

@app.get("/health")
def health():
    return {"status": "ok", "whisper": "base", "rubert": "aniemore-9-emotions"}

@app.get("/metrics")
def metrics():
    return generate_latest()