bitsandbytes for Model Loading

sql_query_generator/generator.py

@@ -2,10 +2,14 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 def load_model(model_name="chatdb/natural-sql-7b"):
     """
-    Loads the SQL generation model and tokenizer from Hugging Face.
+    Loads the SQL generation model with 8-bit precision.
     """
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.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, quantization_config=quantization_config)
     return tokenizer, model
 
 def generate_sql(question, prompt_inputs, tokenizer, model, device="cpu"):