Update app.py

app.py

@@ -11,6 +11,7 @@ from pathlib import Path
 from datetime import datetime
 from typing import List, Dict, Union, Tuple, Optional, Any
 from dataclasses import dataclass, field
+import zipfile
 
 # Machine Learning and Deep Learning Libraries
 import torch
@@ -81,6 +82,24 @@ 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"""