Update app.py

app.py

@@ -130,20 +130,20 @@ class SentenceTransformerRetriever:
         self.doc_embeddings = None
         os.makedirs(cache_dir, exist_ok=True)
         # Initialize model using cached method
-        self.model = self._load_model()
+        self.model = self._load_model(model_name)  # Pass model_name as argument
     
     @st.cache_resource(show_spinner=False)
-    def _load_model(self):
+    def _load_model(_self, _model_name: str):  # Changed to _self and added _model_name
         """Load and cache the sentence transformer model"""
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            model = SentenceTransformer(self.model_name, device="cpu")
-            # Verify model is loaded correctly
-            test_embedding = model.encode("test", convert_to_tensor=True)
-            if not isinstance(test_embedding, torch.Tensor):
-                raise ValueError("Model initialization failed")
-            return model
-    
+        try:
+            with warnings.catch_warnings():
+                warnings.simplefilter("ignore")
+                model = SentenceTransformer(_model_name, device="cpu")
+                # Verify model is loaded correctly
+                test_embedding = model.encode("test", convert_to_tensor=True)
+                if not isinstance(test_embedding, torch.Tensor):
+                    raise ValueError("Model initialization failed")
+                return model
     def get_cache_path(self, data_folder: str = None) -> str:
         return os.path.join(self.cache_dir, self.cache_file)
     
@@ -162,9 +162,9 @@ class SentenceTransformerRetriever:
     
     @log_function
     @st.cache_data
-    def load_cache(self, data_folder: str = None) -> Optional[Dict]:
+    def load_cache(_self, _data_folder: str = None) -> Optional[Dict]:  # Changed to _self and _data_folder
         try:
-            cache_path = self.get_cache_path()
+            cache_path = _self.get_cache_path()
             if os.path.exists(cache_path):
                 with open(cache_path, 'rb') as f:
                     logging.info(f"Loading cache from: {cache_path}")
@@ -207,7 +207,7 @@ class SentenceTransformerRetriever:
         return indices.cpu(), scores.cpu()
 
 class RAGPipeline:
-    def __init__(self, data_folder: str, k: int = 5):
+     def __init__(self, data_folder: str, k: int = 5):
         self.data_folder = data_folder
         self.k = k
         self.retriever = SentenceTransformerRetriever()
@@ -218,20 +218,20 @@ class RAGPipeline:
         self._initialize_model()
 
     @st.cache_resource(show_spinner=False)
-    def _initialize_model(self):
+    def _initialize_model(_self):  # Changed to _self
         """Initialize the model with proper error handling and verification"""
         try:
-            os.makedirs(os.path.dirname(self.model_path), exist_ok=True)
+            os.makedirs(os.path.dirname(_self.model_path), exist_ok=True)
             
-            if not os.path.exists(self.model_path):
+            if not os.path.exists(_self.model_path):
                 direct_url = "https://huggingface.co/TheBloke/Mistral-7B-v0.1-GGUF/resolve/main/mistral-7b-v0.1.Q4_K_M.gguf"
-                download_file_with_progress(direct_url, self.model_path)
+                download_file_with_progress(direct_url, _self.model_path)
 
-            if not os.path.exists(self.model_path):
-                raise FileNotFoundError(f"Model file {self.model_path} not found after download attempts")
+            if not os.path.exists(_self.model_path):
+                raise FileNotFoundError(f"Model file {_self.model_path} not found after download attempts")
             
-            if os.path.getsize(self.model_path) < 1000000:  # Less than 1MB
-                os.remove(self.model_path)
+            if os.path.getsize(_self.model_path) < 1000000:  # Less than 1MB
+                os.remove(_self.model_path)
                 raise ValueError("Downloaded model file is too small, likely corrupted")
             
             llm_config = {
@@ -242,7 +242,7 @@ class RAGPipeline:
                 "verbose": False
             }
         
-            self.llm = Llama(model_path=self.model_path, **llm_config)
+            _self.llm = Llama(model_path=_self.model_path, **llm_config)
             st.success("Model loaded successfully!")
 
         except Exception as e: