Update app.py

app.py

@@ -82,47 +82,23 @@ os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Define the path for the zipped model
-model_zip_path = "./checkpoint-500.zip"
-extracted_model_dir = "./checkpoint-500"
-
-# Unzip the file if it’s not already extracted
-if not os.path.exists(extracted_model_dir):
-    with zipfile.ZipFile(model_zip_path, 'r') as zip_ref:
-        zip_ref.extractall(extracted_model_dir)
-
-# Load the model from the extracted directory
-self.model = AutoModelForCausalLM.from_pretrained(
-    extracted_model_dir,
-    device_map="auto",
-    load_in_8bit=True,
-    torch_dtype=torch.float16,
-    low_cpu_mem_usage=True
-)
+
 
 class ModelManager:
     """Handles model loading and resource management"""
     
     @staticmethod
-    def verify_model_path(checkpoint_path: str) -> str:
-        """Verify and return valid model path"""
-        if os.path.exists(checkpoint_path):
-            return checkpoint_path
+    def verify_and_extract_model(checkpoint_zip_path: str, extracted_model_dir: str) -> str:
+        """Verify and extract the model if it's not already extracted"""
+        if not os.path.exists(extracted_model_dir):
+            # Unzip the model if it hasn’t been extracted yet
+            with zipfile.ZipFile(checkpoint_zip_path, 'r') as zip_ref:
+                zip_ref.extractall(extracted_model_dir)
+            logger.info(f"Extracted model to: {extracted_model_dir}")
+        else:
+            logger.info(f"Model already extracted: {extracted_model_dir}")
             
-        alternate_paths = [
-            f"{os.getcwd()}/checkpoint-500.zip",
-            "./checkpoint-500.zip",
-            "../checkpoint-500.zip"
-        ]
-        
-        for path in alternate_paths:
-            if os.path.exists(path):
-                return path
-                
-        raise FileNotFoundError(
-            f"Model checkpoint not found in any of these locations: "
-            f"{[checkpoint_path] + alternate_paths}"
-        )
+        return extracted_model_dir
     
     @staticmethod
     def clear_gpu_memory():
@@ -132,8 +108,9 @@ class ModelManager:
             gc.collect()
 
 class PearlyBot:
-    def __init__(self, model_path: str = "./checkpoint-500.zip"):
-        self.setup_model(model_path)
+    def __init__(self, model_zip_path: str = "./checkpoint-500.zip", model_dir: str = "./checkpoint-500"):
+        self.model_dir = ModelManager.verify_and_extract_model(model_zip_path, model_dir)
+        self.setup_model(self.model_dir)
         self.setup_rag()
         self.conversation_history = []
         self.last_interaction_time = time.time()
@@ -145,17 +122,9 @@ class PearlyBot:
             logger.info("Starting model initialization...")
             ModelManager.clear_gpu_memory()
             
-            # Verify model path
-            verified_path = ModelManager.verify_model_path(model_path)
-            logger.info(f"Using model checkpoint from: {verified_path}")
-            
-            # Base model configuration
-            base_model_id = "google/gemma-2b"
-            logger.info(f"Loading base model: {base_model_id}")
-            
             # Load tokenizer
             try:
-                self.tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+                self.tokenizer = AutoTokenizer.from_pretrained(model_path)
                 self.tokenizer.pad_token = self.tokenizer.eos_token
                 logger.info("Tokenizer loaded successfully")
             except Exception as e:
@@ -165,7 +134,7 @@ class PearlyBot:
             # Load model
             try:
                 self.model = AutoModelForCausalLM.from_pretrained(
-                    verified_path,
+                    model_path,
                     device_map="auto",
                     load_in_8bit=True,
                     torch_dtype=torch.float16,