add cpu bitsandbytes

requirements.txt

@@ -1,4 +1,5 @@
 streamlit
 transformers
-bitsandbytes
+bitsandbytes-cpu
+
 accelerate

sql_query_generator/generator.py

@@ -1,16 +1,15 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
 
 def load_model(model_name="chatdb/natural-sql-7b"):
     """
-    Loads the model on CPU and avoids bitsandbytes.
+    Loads the SQL generation model with 8-bit precision.
     """
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        device_map="auto",  # Auto-map to CPU
-        offload_folder="offload",  # Offload to disk
-        low_cpu_mem_usage=True,  # Optimize CPU memory usage
+    quantization_config = BitsAndBytesConfig(
+        load_in_8bit=True,  # Enable 8-bit loading
+        llm_int8_threshold=6.0  # Fine-tune threshold if needed
     )
+    model = AutoModelForCausalLM.from_pretrained(model_name, quantization_config=quantization_config)
     return tokenizer, model
 
 def generate_sql(question, prompt_inputs, tokenizer, model, device="cpu"):