Update app.py

app.py

@@ -135,42 +135,6 @@ def check_environment():
         st.stop()
         return False
 
-@st.cache_resource
-def initialize_model(self):
-    try:
-        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)
-
-        # Verify file exists and has content
-        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)
-            raise ValueError("Downloaded model file is too small, likely corrupted")
-            
-        llm_config = {
-            "n_ctx": 2048,
-            "n_threads": 4,
-            "n_batch": 512,
-            "n_gpu_layers": 0,
-            "verbose": False
-        }
-        
-        return Llama(model_path=self.model_path, **llm_config)
-        
-    except FileNotFoundError as e:
-        logging.error(f"Failed to find or download model file: {str(e)}")
-        raise
-        
-    except ValueError as e:
-        logging.error(f"Model file validation failed: {str(e)}")
-        raise
-        
-    except Exception as e:
-        logging.error(f"Unexpected error during model initialization: {str(e)}")
-        raise RuntimeError(f"Failed to initialize model: {str(e)}") from e
 
 
 class SentenceTransformerRetriever:
@@ -188,6 +152,44 @@ class SentenceTransformerRetriever:
     
     def get_cache_path(self, data_folder: str = None) -> str:
         return os.path.join(self.cache_dir, self.cache_file)
+
+    @st.cache_resource
+    def initialize_model(self):
+        try:
+            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)
+
+            # Verify file exists and has content
+            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)
+                raise ValueError("Downloaded model file is too small, likely corrupted")
+            
+            llm_config = {
+                "n_ctx": 2048,
+                "n_threads": 4,
+                "n_batch": 512,
+                "n_gpu_layers": 0,
+                "verbose": False
+            }
+        
+            return Llama(model_path=self.model_path, **llm_config)
+        
+        except FileNotFoundError as e:
+            logging.error(f"Failed to find or download model file: {str(e)}")
+            raise
+        
+        except ValueError as e:
+            logging.error(f"Model file validation failed: {str(e)}")
+            raise
+        
+        except Exception as e:
+            logging.error(f"Unexpected error during model initialization: {str(e)}")
+            raise RuntimeError(f"Failed to initialize model: {str(e)}") from e
+
     
     @log_function
     def save_cache(self, data_folder: str, cache_data: dict):
@@ -248,6 +250,8 @@ class RAGPipeline:
         self.documents = []
         self.device = torch.device("cpu")
         self.llm = initialize_model()
+
+        self.model_path = "mistral-7b-v0.1.Q4_K_M.gguf"
         
     @log_function
     @st.cache_data