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 (
@@ -43,15 +43,17 @@ class SQLGenerator:
         outputs = self.model.generate(
             inputs["input_ids"],
             max_length=256,
-            do_sample=False,
+            do_sample=True,  # Enable sampling to use temperature
             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
+            temperature=0.7,  # Allow temperature to affect output
             top_k=50          # Consider top k predictions for variability
         )
-        return self.tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
-
+        
+        generated_query = self.tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+        return generated_query  # Return the generated SQL query
+    
     def fetch_shopify_data(self, endpoint):
         headers = {
             'X-Shopify-Access-Token': ACCESS_TOKEN,
@@ -59,9 +61,24 @@ 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
+
+# Example of how to use the SQLGenerator class
+if __name__ == "__main__":
+    sql_generator = SQLGenerator()
+    
+    # Example natural language query
+    natural_language_query = "Show me shirts with red color"
+    
+    # Generate SQL query
+    sql_query = sql_generator.generate_query(natural_language_query)
+    print(f"Generated SQL Query: {sql_query}")
+    
+    # Fetch data from Shopify (example endpoint)
+    shopify_data = sql_generator.fetch_shopify_data("products")
+    print(f"Shopify Data: {shopify_data}")