Update main.py

main.py

@@ -1,5 +1,5 @@
 import fitz
-from fastapi import FastAPI, File, UploadFile, Form
+from fastapi import FastAPI, File, UploadFile, Form, HTTPException
 from fastapi.responses import JSONResponse
 from transformers import pipeline
 from PIL import Image
@@ -12,6 +12,8 @@ import numpy as np
 import json
 import torchaudio
 import torch
+from pydub import AudioSegment
+import speech_recognition as sr
 
 app = FastAPI()
 
@@ -164,34 +166,34 @@ async def transcribe_and_answer(
     file: UploadFile = File(...),
     questions: str = Form(...)
 ):
+    # Check the file format and read it
+    if file.content_type not in ["audio/wav", "audio/mpeg", "audio/mp3"]:
+        raise HTTPException(status_code=400, detail="Unsupported audio format. Please upload a WAV or MP3 file.")
+    
+    # Convert uploaded file to WAV if necessary (for SpeechRecognition compatibility)
+    audio_data = await file.read()
+    audio_file = io.BytesIO(audio_data)
+    
     try:
-        # Step 1: Read and convert the audio file
-        contents = await file.read()
-        audio = AudioSegment.from_file(BytesIO(contents))
-
-        # Step 2: Ensure the audio is mono and resample if needed
-        audio = audio.set_channels(1)  # Convert to mono if it's not already
-        audio = audio.set_frame_rate(16000)  # Resample to 16000 Hz, commonly required by ASR models
-
-        # Step 3: Export to WAV format and load with torchaudio
-        wav_buffer = BytesIO()
-        audio.export(wav_buffer, format="wav")
-        wav_buffer.seek(0)
-
-        # Load audio using torchaudio
-        waveform, sample_rate = torchaudio.load(wav_buffer)
+        # Convert MP3 to WAV if needed
+        if file.content_type == "audio/mpeg" or file.content_type == "audio/mp3":
+            audio = AudioSegment.from_mp3(audio_file)
+            audio_wav = io.BytesIO()
+            audio.export(audio_wav, format="wav")
+            audio_wav.seek(0)
+        else:
+            audio_wav = audio_file
         
-        # Convert waveform to float32 and ensure it's a numpy array
-        waveform_np = waveform.numpy().astype(np.float32)
-
-        # Step 4: Transcribe the audio
-        transcription_result = nlp_speech_to_text(waveform_np)
-        transcription_text = transcription_result['text']
-
-        # Step 5: Parse the JSON-formatted questions
+        # Load audio into SpeechRecognition and transcribe
+        recognizer = sr.Recognizer()
+        with sr.AudioFile(audio_wav) as source:
+            audio = recognizer.record(source)
+            transcription_text = recognizer.recognize_google(audio)
+        
+        # Parse the JSON-formatted questions
         questions_dict = json.loads(questions)
 
-        # Step 6: Answer each question using the transcribed text
+        # Answer each question based on the transcription text
         answers_dict = {}
         for key, question in questions_dict.items():
             QA_input = {
@@ -202,14 +204,14 @@ async def transcribe_and_answer(
             result = nlp_qa_v3(QA_input)
             answers_dict[key] = result['answer']
 
-        # Step 7: Return transcription + answers
+        # Return transcription + answers
         return {
             "transcription": transcription_text,
             "answers": answers_dict
         }
 
     except Exception as e:
-        return JSONResponse(content={"error": f"Error processing audio or answering questions: {str(e)}"}, status_code=500)
+        raise HTTPException(status_code=500, detail=f"Error during transcription or question answering: {str(e)}")
 
 # Set up CORS middleware
 origins = ["*"]  # or specify your list of allowed origins