app.py

import shutil
import tempfile
import os

from fastapi import FastAPI, File, UploadFile, HTTPException
from fastapi.responses import JSONResponse
from transformers import pipeline

app = FastAPI(
    title="Voice Safety Detection API",
    description="Upload an audio file. The API transcribes it (Whisper) then classifies the speech as Safe or Danger.",
    version="1.0.0",
)

# ---- تحميل الموديلات مرة واحدة عند بدء السيرفر ----
whisper = None
classifier = None

label_map = {
    "LABEL_0": "Safe",
    "LABEL_1": "Danger",
}


@app.on_event("startup")
def load_models():
    global whisper, classifier
    whisper = pipeline(
        "automatic-speech-recognition",
        model="openai/whisper-base",
    )
    classifier = pipeline(
        "text-classification",
        model="MennatullahHany/Abert",
    )


@app.get("/")
def root():
    return {"status": "ok", "message": "Voice Safety Detection API is running."}


@app.get("/health")
def health():
    return {"status": "healthy"}


@app.post("/predict")
async def predict(audio: UploadFile = File(...)):
    if whisper is None or classifier is None:
        raise HTTPException(status_code=503, detail="Models are still loading, try again shortly.")

    # نحفظ الملف المرفوع مؤقتًا على القرص لأن pipeline يحتاج مسار ملف
    suffix = os.path.splitext(audio.filename or "")[1] or ".wav"
    tmp_path = None
    try:
        with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
            shutil.copyfileobj(audio.file, tmp)
            tmp_path = tmp.name

        transcription = whisper(tmp_path)["text"]
        result = classifier(transcription)[0]
        label = label_map.get(result["label"], result["label"])

        return JSONResponse(
            content={
                "text": transcription,
                "result": label,
                "confidence": round(float(result["score"]), 4),
            }
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.remove(tmp_path)