Update app.py

app.py

@@ -1,6 +1,4 @@
 import os
-os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf-cache"  # Important for Docker
-
 from fastapi import FastAPI, UploadFile, File
 from fastapi.middleware.cors import CORSMiddleware
 import torchaudio
@@ -8,9 +6,13 @@ import torch
 from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
 import io
 
+# DO NOT set the cache directory here anymore.
+# Let the Dockerfile's ENV variables handle it.
+# REMOVED: os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf-cache"
+
 app = FastAPI()
 
-# Allow all origins (for Flutter)
+# Allow all origins
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -18,9 +20,19 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# Load model
-processor = Wav2Vec2Processor.from_pretrained("Mustafaa4a/ASR-Somali")
-model = Wav2Vec2ForCTC.from_pretrained("Mustafaa4a/ASR-Somali")
+# --- Model Loading ---
+# This will now use the cache path set by the Dockerfile's ENV variables (/app/hf-cache)
+print("Loading model and processor...")
+try:
+    processor = Wav2Vec2Processor.from_pretrained("Mustafaa4a/ASR-Somali")
+    model = Wav2Vec2ForCTC.from_pretrained("Mustafaa4a/ASR-Somali")
+    print("Model and processor loaded successfully.")
+except Exception as e:
+    print(f"FATAL: Could not load model. Error: {e}")
+    # In a real app, you might want to exit or handle this gracefully
+    processor = None
+    model = None
+
 
 @app.get("/")
 async def root():
@@ -28,20 +40,26 @@ async def root():
 
 @app.post("/transcribe")
 async def transcribe(file: UploadFile = File(...)):
-    audio_bytes = await file.read()
-    audio_stream = io.BytesIO(audio_bytes)
-    
-    waveform, sample_rate = torchaudio.load(audio_stream)
-
-    if sample_rate != 16000:
-        resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
-        waveform = resampler(waveform)
-
-    inputs = processor(waveform.squeeze(), sampling_rate=16000, return_tensors="pt")
-    
-    with torch.no_grad():
-        logits = model(**inputs).logits
-
-    predicted_ids = torch.argmax(logits, dim=-1)
-    transcription = processor.decode(predicted_ids[0])
-    return {"transcription": transcription}
+    if not model or not processor:
+        return {"error": "Model is not loaded, please check server logs for errors."}
+
+    try:
+        audio_bytes = await file.read()
+        audio_stream = io.BytesIO(audio_bytes)
+
+        waveform, sample_rate = torchaudio.load(audio_stream)
+
+        if sample_rate != 16000:
+            resampler = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
+            waveform = resampler(waveform)
+
+        inputs = processor(waveform.squeeze(), sampling_rate=16000, return_tensors="pt")
+
+        with torch.no_grad():
+            logits = model(**inputs).logits
+
+        predicted_ids = torch.argmax(logits, dim=-1)
+        transcription = processor.decode(predicted_ids[0])
+        return {"transcription": transcription}
+    except Exception as e:
+        return {"error": f"An error occurred during transcription: {str(e)}"}