Update app.py

app.py

@@ -6,6 +6,7 @@ import logging
 import sys
 import os
 from accelerate import infer_auto_device_map, init_empty_weights
+import bitsandbytes as bnb
 
 # Configure logging
 logging.basicConfig(
@@ -45,7 +46,7 @@ try:
     tokenizer.pad_token = tokenizer.eos_token
     logger.info("Tokenizer loaded successfully")
 
-    # Load model with device map
+    # Load model with quantization and device map
     logger.info("Loading model...")
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
@@ -53,7 +54,12 @@ try:
         trust_remote_code=True,
         token=hf_token,
         device_map="auto",
-        load_in_8bit=True
+        quantization_config={
+            "load_in_4bit": True,
+            "bnb_4bit_compute_dtype": torch.float16,
+            "bnb_4bit_use_double_quant": True,
+            "bnb_4bit_quant_type": "nf4"
+        }
     )
     logger.info("Model loaded successfully")