Update sql_generator.py

sql_generator.py

@@ -8,7 +8,7 @@ class SQLGenerator:
         self.model_name = "premai-io/prem-1B-SQL"
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
         self.model = AutoModelForCausalLM.from_pretrained(self.model_name)
-        
+
     def generate_query(self, natural_language_query):
         schema_info = """
         CREATE TABLE products (
@@ -34,23 +34,24 @@ class SQLGenerator:
         """
         
         prompt = f"""### Task: Generate a SQL query to answer the following question.
-        ### Database Schema:
-        {schema_info}
-        ### Question: {natural_language_query}
-        ### SQL Query:"""
-        
+### Database Schema:
+{schema_info}
+### Question: {natural_language_query}
+### SQL Query:"""
+
+        inputs = self.tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(self.model.device)
         outputs = self.model.generate(
-        inputs["input_ids"],
-        max_length=256,
-        do_sample=False,
-        num_return_sequences=1,
-        eos_token_id=self.tokenizer.eos_token_id,
-        pad_token_id=self.tokenizer.pad_token_id,
-        temperature=0.7,  # Adjust temperature for more creative output
-        top_k=50          # Consider top k predictions for variability
-    )
+            inputs["input_ids"],
+            max_length=256,
+            do_sample=False,
+            num_return_sequences=1,
+            eos_token_id=self.tokenizer.eos_token_id,
+            pad_token_id=self.tokenizer.pad_token_id,
+            temperature=0.7,  # Adjust temperature for more creative output
+            top_k=50          # Consider top k predictions for variability
+        )
         return self.tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
-    
+
     def fetch_shopify_data(self, endpoint):
         headers = {
             'X-Shopify-Access-Token': ACCESS_TOKEN,
@@ -58,9 +59,9 @@ class SQLGenerator:
         }
         url = f"https://{SHOP_NAME}/admin/api/2023-10/{endpoint}.json"
         response = requests.get(url, headers=headers)
-        
+
         if response.status_code == 200:
             return response.json()
         else:
             print(f"Error fetching {endpoint}: {response.status_code} - {response.text}")
-            return None
+            return None