Update fast_api.py

fast_api.py

@@ -8,16 +8,32 @@ import pickle
 import numpy as np
 from gensim.models import KeyedVectors
 
+def load_whisper_model(model_path, device='cpu'):
+    # Load model architecture
+    model = whisper.model.Whisper(
+        whisper.model.ModelDimensions(
+            n_mels=80,
+            n_audio_ctx=1500,
+            n_audio_state=384,
+            n_audio_head=6,
+            n_audio_layer=4,
+            n_vocab=51865,
+            n_text_ctx=448,
+            n_text_state=384,
+            n_text_head=6,
+            n_text_layer=4
+        )
+    )
+    
+    # Load state dict
+    state_dict = torch.load(model_path, map_location=device)
+    model.load_state_dict(state_dict)
+    model.eval()
+    return model
+    
 # Load the saved Word2Vec model
 word2vec_model = KeyedVectors.load("word2vec-google-news-300.model")
-
-model = whisper.load_model("tiny")
-
-# Load the saved state dictionary
-model_state = torch.load("whisper_tiny_model.pt")
-
-# Load the state dictionary into the model
-model.load_state_dict(model_state)
+model = load_whisper_model("whisper_tiny_model.pt")
 
 def load_model(pickle_file_path: str):
     """Load a model from a pickle file."""
@@ -25,12 +41,10 @@ def load_model(pickle_file_path: str):
         model = pickle.load(file)
     return model
 
-
 pronunciation_fluency_model = load_model("pronunciation_fluency_v2.pkl")
 
 app = FastAPI()
 
-
 def transcribe(audio_file_path: str, model):
     # Load audio and run inference
     result = model.transcribe(audio_file_path)