Update app.py

app.py

@@ -1,17 +1,26 @@
 from fastapi import FastAPI
 import uvicorn
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 model_name = "TheBloke/Wizard-Vicuna-13B-Uncensored-HF"
 
+# Configure 4-bit quantization
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,               # Enable 4-bit quantization
+    bnb_4bit_quant_type="nf4",       # Use 4-bit NormalFloat (optimal)
+    bnb_4bit_compute_dtype="float16", # Faster computation with float16
+    bnb_4bit_use_double_quant=True   # Extra compression
+)
+
+# Load model with quantization
 model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    device_map="auto",            # Auto-distribute layers across CPU/GPU
-    low_cpu_mem_usage=True,       # Reduces CPU RAM during loading
-    torch_dtype="auto",           # Automatically select dtype (float16/32)
-    offload_folder="offload",
-    trust_remote_code=True
+    model_name,  # Example model
+    quantization_config=bnb_config,
+    device_map="auto",            # Auto-distribute across GPU/CPU
+    trust_remote_code=True        # Required for some models
 )
+
+# load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 tokenizer.pad_token = tokenizer.eos_token