Update app.py

app.py

@@ -261,46 +261,52 @@ class Qwen25SmallLLM(Runnable):
     """LLM class that properly inherits from Runnable for LangChain compatibility"""
     
     def __init__(self, model_path: str = "Qwen/Qwen2.5-3B-Instruct", use_4bit: bool = True):
-        super().__init__()
-        logger.info(f"Loading model: {model_path} (use_4bit={use_4bit})")
-        start_Loading_Model_time = time.perf_counter()
+    super().__init__()
+    logger.info(f"Loading model: {model_path} (use_4bit={use_4bit})")
+    start_Loading_Model_time = time.perf_counter()
+    current_time = datetime.now()
 
-        try:
-            # Load tokenizer
-            self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-
-            if use_4bit:
-                quant_config = BitsAndBytesConfig(
-                    load_in_4bit=True,
-                    bnb_4bit_compute_dtype=torch.float16,
-                    bnb_4bit_use_double_quant=True,
-                    bnb_4bit_quant_type="nf4", 
-                    llm_int8_threshold=0.0,
-                    llm_int8_skip_modules=["lm_head"]
-                )
+    try:
+        # Load tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+        if use_4bit:
+            quant_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_use_double_quant=True,
+                bnb_4bit_quant_type="nf4", 
+                llm_int8_threshold=0.0,
+                llm_int8_skip_modules=["lm_head"]
+            )
 
-                # Try quantized load with updated dtype parameter
-                self.model = AutoModelForCausalLM.from_pretrained(
-                    model_path,
-                    quantization_config=quant_config,
-                    device_map="auto",
-                    dtype=torch.bfloat16,
-                    trust_remote_code=True,
-                    low_cpu_mem_usage=True
-                )
-            else:
-                self._load_fallback_model(model_path)
+            # Try quantized load with updated dtype parameter
+            self.model = AutoModelForCausalLM.from_pretrained(
+                model_path,
+                quantization_config=quant_config,
+                device_map="auto",
+                dtype=torch.bfloat16,
+                trust_remote_code=True,
+                low_cpu_mem_usage=True
+            )
+        else:
+            self._load_fallback_model(model_path)
+            
+        # Success path - log timing
         end_Loading_Model_time = time.perf_counter()
         Loading_Model_time = end_Loading_Model_time - start_Loading_Model_time
-
         log_metric(f"Model Load time: {Loading_Model_time:0.4f} seconds. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
-        except Exception as e:
-            logger.warning(f"Quantized load failed, falling back: {e}")
-            self._load_fallback_model(model_path)
+        
+    except Exception as e:
+        logger.warning(f"Quantized load failed, falling back: {e}")
+        self._load_fallback_model(model_path)
+        end_Loading_Model_time = time.perf_counter()
+        Loading_Model_time = end_Loading_Model_time - start_Loading_Model_time
+        log_metric(f"Model Load time (fallback): {Loading_Model_time:0.4f} seconds. Timestamp: {current_time:%Y-%m-%d %H:%M:%S}")
 
-        # Ensure pad token
-        if self.tokenizer.pad_token is None:
-            self.tokenizer.pad_token = self.tokenizer.eos_token
+    # Ensure pad token
+    if self.tokenizer.pad_token is None:
+        self.tokenizer.pad_token = self.tokenizer.eos_token
 
     def _load_fallback_model(self, model_path: str):
         """Fallback if quantization fails."""