Update app/models.py

app/models.py

@@ -31,21 +31,20 @@ class DataLocation(BaseModel):
         """
         if not os.path.exists(self.local_path):
             if self.cloud_uri is not None:
-                logger.warning(f"Downloading model from cloud URI: {self.cloud_uri}")
-                # Implement cloud download logic here if needed
-            else:
-                logger.info(f"Downloading model from Hugging Face to: {self.local_path}")
+                logger.warning(f"Downloading model from Hugging Face: {self.cloud_uri}")
                 # Download from Hugging Face
                 tokenizer = AutoTokenizer.from_pretrained(
-                    self.cloud_uri or self.local_path, use_auth_token=AUTH_TOKEN
+                    self.cloud_uri, use_auth_token=AUTH_TOKEN
                 )
                 model = AlbertForQuestionAnswering.from_pretrained(
-                    self.cloud_uri or self.local_path, use_auth_token=AUTH_TOKEN
+                    self.cloud_uri, use_auth_token=AUTH_TOKEN
                 )
                 # Save the model and tokenizer locally
                 tokenizer.save_pretrained(self.local_path)
                 model.save_pretrained(self.local_path)
                 logger.info(f"Model saved to: {self.local_path}")
+            else:
+                raise ValueError(f"Model not found locally and no cloud URI provided: {self.local_path}")
         return self.local_path
 
 # Define the model location
@@ -64,16 +63,16 @@ class QAModel:
         self.tokenizer = None
         self.model = None
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        self.load_model()
+        self._load_model()  # Call the method to load the model and tokenizer
 
-    def load_model(self):
+    def _load_model(self):
         """
         Load the tokenizer and model.
         """
         # Ensure the model is downloaded
         model_path = MODEL_LOCATION.exists_or_download()
 
-        # Load the tokenizer and model
+        # Load the tokenizer and model from the local path
         self.tokenizer = AutoTokenizer.from_pretrained(model_path)
         self.model = AlbertForQuestionAnswering.from_pretrained(model_path).to(self.device)
         logger.info(f"Loaded QA model: {self.model_name}")
@@ -114,7 +113,7 @@ def load_qa_pipeline():
     Load the QA model and tokenizer.
     """
     global qa_model
-    qa_model = QAModel()
+    qa_model = QAModel()  # This will automatically call `_load_model` during initialization
     return qa_model
 
 def inference_qa(qa_pipeline, context: str, question: str):