CPU compatible quantization

sql_query_generator/generator.py

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