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

app = FastAPI()

# โหลดโมเดล (ใช้ CPU เท่านั้น)
classifier = pipeline("audio-classification", 
                      model="AmeerHesham/distilhubert-finetuned-baby_cry", 
                      device=-1)

@app.post("/predict")
async def predict_audio(file: UploadFile = File(...)):
    try:
        contents = await file.read()
        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
            tmp.write(contents)
            tmp_path = tmp.name

        results = classifier(tmp_path)
        return JSONResponse(content={"results": results})
    except Exception as e:
        return JSONResponse(content={"error": str(e)}, status_code=500)

if __name__ == "__main__":
    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)