Update app.py

app.py

@@ -219,17 +219,12 @@ class Qwen25SmallLLM:
             if use_4bit:
                 quant_config = BitsAndBytesConfig(
                     load_in_4bit=True,
-                    bnb_4bit_compute_dtype=torch.bfloat16,
+                    bnb_4bit_compute_dtype=torch.float16,
                     bnb_4bit_use_double_quant=True,
-                    bnb_4bit_quant_type="nf4"
-                )
-                logger.info("Using 4-bit quantization with BitsAndBytes")
-            else:
-                quant_config = BitsAndBytesConfig(
-                    load_in_8bit=True,
-                    llm_int8_enable_fp32_cpu_offload=True
-                )
-                logger.info("Using 8-bit quantization with BitsAndBytes")
+                    bnb_4bit_quant_type="nf4", 
+                    llm_int8_threshold=0.0,
+                    llm_int8_skip_modules=["lm_head"]
+                    )
 
             # Try quantized load
             self.model = AutoModelForCausalLM.from_pretrained(