feature to save model locally

app/models.py

@@ -1,12 +1,66 @@
+import os
+import logging
+from pathlib import Path
+from typing import Optional
+from pydantic import BaseModel
 from transformers import AutoTokenizer, AlbertForQuestionAnswering
 import torch
 
+# Set up logging
+logger = logging.getLogger(__name__)
+
+# Define the model directory
+MODEL_DIR = Path(__file__).parent.parent / "_models"
+MODEL_DIR.mkdir(parents=True, exist_ok=True)  # Create the directory if it doesn't exist
+
+# Hugging Face authentication token (from environment variable)
+AUTH_TOKEN = os.getenv("auth_token")
+if not AUTH_TOKEN:
+    raise ValueError("Hugging Face auth_token environment variable is not set.")
+
+class DataLocation(BaseModel):
+    """
+    Represents the location of a model (local path and optional cloud URI).
+    """
+    local_path: str
+    cloud_uri: Optional[str] = None
+
+    def exists_or_download(self):
+        """
+        Check if the model exists locally. If not, download it from Hugging Face.
+        """
+        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}")
+                # Download from Hugging Face
+                tokenizer = AutoTokenizer.from_pretrained(
+                    self.cloud_uri or self.local_path, use_auth_token=AUTH_TOKEN
+                )
+                model = AlbertForQuestionAnswering.from_pretrained(
+                    self.cloud_uri or self.local_path, 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}")
+        return self.local_path
+
+# Define the model location
+MODEL_NAME = "twmkn9/albert-base-v2-squad2"
+MODEL_LOCATION = DataLocation(
+    local_path=str(MODEL_DIR / MODEL_NAME.replace("/", "-")),
+    cloud_uri=MODEL_NAME,  # Hugging Face model ID
+)
+
 class QAModel:
-    def __init__(self, model_name: str = "Gowtham122/albertqa"):
+    def __init__(self):
         """
         Initialize the QA model and tokenizer.
         """
-        self.model_name = model_name
+        self.model_name = MODEL_NAME
         self.tokenizer = None
         self.model = None
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -16,8 +70,12 @@ class QAModel:
         """
         Load the tokenizer and model.
         """
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
-        self.model = AlbertForQuestionAnswering.from_pretrained(self.model_name).to(self.device)
+        # Ensure the model is downloaded
+        model_path = MODEL_LOCATION.exists_or_download()
+
+        # Load the tokenizer and model
+        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}")
 
     def inference_qa(self, context: str, question: str):
@@ -51,12 +109,12 @@ class QAModel:
 # Global instance of the QA model
 qa_model = QAModel()
 
-def load_qa_pipeline(model_name: str = "Gowtham122/albertqa"):
+def load_qa_pipeline():
     """
     Load the QA model and tokenizer.
     """
     global qa_model
-    qa_model = QAModel(model_name)
+    qa_model = QAModel()
     return qa_model
 
 def inference_qa(qa_pipeline, context: str, question: str):