from fastapi import FastAPI, File, UploadFile
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
import numpy as np

app = FastAPI()
mapper = ["angry", "disgust", "fear", "happy", "neutral", "other", "sad", "surprised", "unknown"]
inference_pipeline = pipeline(
    task=Tasks.emotion_recognition,
    model="iic/emotion2vec_base_finetuned",
    model_revision="v2.0.4"
)

CHUNK_SIZE = 44100  # Adjust the chunk size as needed

@app.post("/emotion_recognition")
async def emotion_recognition(file: UploadFile = File(...)):
    rec_result = []
    while True:
        chunk = await file.read(CHUNK_SIZE)
        if not chunk:
            break
        if len(chunk) % 44100 != 0:
            # Skip the last chunk if it's not a multiple of the expected element size
            continue
        chunk_result = inference_pipeline(
            chunk, output_dir="./outputs", granularity="utterance", extract_embedding=False
        )
        print(chunk_result)
        rec_result.extend(chunk_result)
    max_emotion_score = np.argmax(rec_result[0]["scores"])
    return {"emotion": mapper[max_emotion_score]}