Add audio transcription feature using Whisper model and update Dockerfile

Dockerfile

@@ -1,5 +1,6 @@
 FROM python:3.9
 
+RUN apt-get update && apt-get install -y ffmpeg && rm -rf /var/lib/apt/lists/*
 RUN useradd -m -u 1000 user
 USER user
 ENV PATH="/home/user/.local/bin:$PATH"

main.py

@@ -1,11 +1,12 @@
-from fastapi import FastAPI, Body
+from fastapi import FastAPI, Body, UploadFile, File
 import torch
 import spacy
 import os
 from pathlib import Path
 from fastapi.middleware.cors import CORSMiddleware
-from transformers import AutoTokenizer, AutoModelForTokenClassification, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForTokenClassification, AutoModelForSequenceClassification, AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
 from pydantic import BaseModel
+import tempfile
 
 # Define input model
 
@@ -99,6 +100,26 @@ nlp = spacy.load('fr_core_news_lg')
 
 # Set device (CPU or GPU)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+
+# Load Whisper model and processor
+model_id = "openai/whisper-large-v3-turbo"
+whisper_model = AutoModelForSpeechSeq2Seq.from_pretrained(
+    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
+)
+whisper_model.to(device)
+
+whisper_processor = AutoProcessor.from_pretrained(model_id)
+
+whisper_pipe = pipeline(
+    "automatic-speech-recognition",
+    model=whisper_model,
+    tokenizer=whisper_processor.tokenizer,
+    feature_extractor=whisper_processor.feature_extractor,
+    torch_dtype=torch_dtype,
+    device=device,
+)
+
 ner_model = ner_model.to(device)
 type_model = type_model.to(device)
 
@@ -200,3 +221,13 @@ async def analyze_text(input_data: TextInput):
         "confidence": confidence,
         "entities": filtered_entities
     }
+
+
+@app.post("/transcribe/")
+async def transcribe_audio(file: UploadFile = File(...)):
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as tmp:
+        tmp.write(await file.read())
+        tmp_path = tmp.name
+
+    result = whisper_pipe(tmp_path)
+    return {"transcription": result["text"]}

requirements.txt

@@ -6,3 +6,5 @@ pydantic==2.9.2
 safetensors==0.4.5
 spacy==3.7.2
 fr-core-news-lg @ https://github.com/explosion/spacy-models/releases/download/fr_core_news_lg-3.7.0/fr_core_news_lg-3.7.0-py3-none-any.whl
+torchaudio
+datasets[audio]