Update app.py

app.py

@@ -1,100 +1,86 @@
-# app.py
-from flask import Flask, request, jsonify
-from flask_cors import CORS
+import os
+import tempfile
 import torch
 import torchaudio
-from transformers import Wav2Vec2ForCTC, Wav2Vec2Tokenizer
-import tempfile
-import os
 import librosa
 import numpy as np
-
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from transformers import Wav2Vec2ForCTC, Wav2Vec2Tokenizer
 from huggingface_hub import login
 
-
 # Handle authentication if token is provided
 hf_token = os.getenv("HF_TOKEN")
 if hf_token:
     login(token=hf_token)
 
-app = Flask(__name__)
-CORS(app)  # Enable CORS for all routes
+# Initialize FastAPI app
+app = FastAPI()
+
+# Enable CORS for all origins (you can restrict this as needed)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Allows all origins; adjust for production
+    allow_credentials=True,
+    allow_methods=["*"],  # Allows all methods
+    allow_headers=["*"],  # Allows all headers
+)
 
 # Load the model and tokenizer
-MODEL_NAME = "STT-Darija-ORG/wav2vec2-xlsr-300m-darija-no-augmentation-v2"  # You can replace with your fine-tuned model
+MODEL_NAME = "STT-Darija-ORG/wav2vec2-xlsr-300m-darija-no-augmentation-v2"
 tokenizer = Wav2Vec2Tokenizer.from_pretrained(MODEL_NAME)
 model = Wav2Vec2ForCTC.from_pretrained(MODEL_NAME)
 
 def preprocess_audio(audio_path):
     """Preprocess audio file for the model"""
-    # Load audio file
     speech, sample_rate = librosa.load(audio_path, sr=16000)
-    
-    # Ensure audio is the right format
     if len(speech.shape) > 1:
         speech = speech.mean(axis=1)  # Convert to mono if stereo
-    
     return speech, sample_rate
 
-@app.route('/transcribe', methods=['POST'])
-def transcribe_audio():
+@app.post("/transcribe")
+def transcribe_audio(file: UploadFile = File(...)):
     try:
-        # Check if audio file is present
-        if 'audio' not in request.files:
-            return jsonify({'error': 'No audio file provided'}), 400
-        
-        audio_file = request.files['audio']
-        
-        if audio_file.filename == '':
-            return jsonify({'error': 'No file selected'}), 400
-        
         # Save the uploaded file temporarily
         with tempfile.NamedTemporaryFile(delete=False, suffix='.wav') as temp_file:
-            audio_file.save(temp_file.name)
+            temp_file.write(file.file.read())
             temp_path = temp_file.name
         
-        try:
-            # Preprocess the audio
-            speech, sample_rate = preprocess_audio(temp_path)
-            
-            # Tokenize the audio
-            input_values = tokenizer(speech, return_tensors="pt", sampling_rate=sample_rate).input_values
-            
-            # Perform inference
-            with torch.no_grad():
-                logits = model(input_values).logits
-            
-            # Decode the predictions
-            predicted_ids = torch.argmax(logits, dim=-1)
-            transcription = tokenizer.decode(predicted_ids[0])
-            
-            # Clean up the temporary file
-            os.unlink(temp_path)
-            
-            return jsonify({'transcription': transcription})
-            
-        except Exception as e:
-            # Clean up the temporary file in case of error
-            if os.path.exists(temp_path):
-                os.unlink(temp_path)
-            raise e
-            
+        # Preprocess the audio
+        speech, sample_rate = preprocess_audio(temp_path)
+        
+        # Tokenize the audio
+        input_values = tokenizer(speech, return_tensors="pt", sampling_rate=sample_rate).input_values
+        
+        # Perform inference
+        with torch.no_grad():
+            logits = model(input_values).logits
+        
+        # Decode the predictions
+        predicted_ids = torch.argmax(logits, dim=-1)
+        transcription = tokenizer.decode(predicted_ids[0])
+        
+        # Clean up the temporary file
+        os.unlink(temp_path)
+        
+        return {"transcription": transcription}
+    
     except Exception as e:
-        return jsonify({'error': f'Transcription failed: {str(e)}'}), 500
+        # Clean up the temporary file in case of error
+        if os.path.exists(temp_path):
+            os.unlink(temp_path)
+        raise HTTPException(status_code=500, detail=f"Transcription failed: {str(e)}")
 
-@app.route('/health', methods=['GET'])
+@app.get("/health")
 def health_check():
-    return jsonify({'status': 'healthy', 'model': MODEL_NAME})
+    return {"status": "healthy", "model": MODEL_NAME}
 
-@app.route('/', methods=['GET'])
+@app.get("/")
 def home():
-    return jsonify({
-        'message': 'Speech Recognition API',
-        'endpoints': {
-            '/transcribe': 'POST - Upload audio file for transcription',
-            '/health': 'GET - Check API health'
+    return {
+        "message": "Speech Recognition API",
+        "endpoints": {
+            "/transcribe": "POST - Upload audio file for transcription",
+            "/health": "GET - Check API health"
         }
-    })
-
-if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860)  # HuggingFace Spaces uses port 7860
+    }