Create db_query.py

db_query.py

@@ -0,0 +1,131 @@
+import gradio as gr
+import requests
+from transformers import AutoModelForCausalLM, AutoTokenizer, DistilBertTokenizer, DistilBertForSequenceClassification
+
+# Shopify API credentials
+API_KEY = 'c7c07dac8f8eaebd316ac4aa16350472'
+API_SECRET = '046b2c0bfb7879db2bba86f24a281ec2'
+ACCESS_TOKEN = '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'
+SHOP_NAME = '6znwwf-77.myshopify.com'
+
+# Load the SQL generation model
+model_name = "premai-io/prem-1B-SQL"
+tokenizer_sql = AutoTokenizer.from_pretrained(model_name)
+model_sql = AutoModelForCausalLM.from_pretrained(model_name)
+
+# Load the DistilGPT-2 model and tokenizer for processing results
+distil_gpt2_model_name = "distilgpt2"
+tokenizer_gpt2 = AutoTokenizer.from_pretrained(distil_gpt2_model_name)
+model_gpt2 = AutoModelForCausalLM.from_pretrained(distil_gpt2_model_name)
+
+def fetch_data_from_shopify(endpoint):
+    """Fetch data from Shopify API."""
+    headers = {
+        'X-Shopify-Access-Token': ACCESS_TOKEN,
+        'Content-Type': 'application/json'
+    }
+    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
+
+def generate_sql(natural_language_query):
+    """Generate SQL query from natural language and retrieve data from Shopify."""
+    
+    # Define your schema information
+    schema_info = """
+    CREATE TABLE products (
+      id DECIMAL(8,2) PRIMARY KEY,
+      title VARCHAR(255),
+      body_html VARCHAR(255),
+      vendor VARCHAR(255),
+      product_type VARCHAR(255),
+      created_at VARCHAR(255),
+      handle VARCHAR(255),
+      updated_at DATE,
+      published_at VARCHAR(255),
+      template_suffix VARCHAR(255),
+      published_scope VARCHAR(255),
+      tags VARCHAR(255),
+      status VARCHAR(255),
+      admin_graphql_api_id DECIMAL(8,2),
+      variants VARCHAR(255),
+      options VARCHAR(255),
+      images VARCHAR(255),
+      image VARCHAR(255)
+    );
+    """
+    
+    # Construct the prompt for SQL generation
+    prompt = f"""### Task: Generate a SQL query to answer the following question.
+### Database Schema:
+{schema_info}
+### Question: {natural_language_query}
+### SQL Query:"""
+    # Tokenize and generate SQL query
+    inputs = tokenizer_sql(prompt, return_tensors="pt", add_special_tokens=False).to(model_sql.device)
+    outputs = model_sql.generate(
+        inputs["input_ids"],
+        max_length=256,  # Keep it short to speed up generation
+        do_sample=False,  # Use greedy search for faster results
+        num_return_sequences=1,
+        eos_token_id=tokenizer_sql.eos_token_id,
+        pad_token_id=tokenizer_sql.pad_token_id
+    )
+
+    # Decode and clean up the response
+    generated_query = tokenizer_sql.decode(outputs[0], skip_special_tokens=True)
+    sql_query = generated_query.strip()  # Clean the output to get the SQL query
+
+
+    # Use the generated SQL query to determine the endpoint and parameters
+    endpoint = "products"  # This can be dynamic based on your SQL parsing logic
+    products = fetch_data_from_shopify(endpoint)
+
+    # Prepare results based on the Shopify data
+    if products and 'products' in products:
+        results = products['products']  # Assuming you want the product data
+
+        # Process results through DistilGPT-2 model
+        result_text = "\n".join([f"Title: {product['title']}, Vendor: {product['vendor']}" for product in results])
+        
+        # Prepare prompt for DistilGPT-2
+        gpt2_prompt = f"Here are the products based on your query:\n{result_text}\n\nProvide a summary of these products."
+
+        # Tokenize and generate response
+        inputs_gpt2 = tokenizer_gpt2(gpt2_prompt, return_tensors="pt", add_special_tokens=False).to(model_gpt2.device)
+        outputs_gpt2 = model_gpt2.generate(
+            inputs_gpt2["input_ids"],
+            max_length=150,
+            temperature=0.7,
+            do_sample=True,
+            top_p=0.95,
+            num_return_sequences=1,
+            eos_token_id=tokenizer_gpt2.eos_token_id,
+            pad_token_id=tokenizer_gpt2.pad_token_id
+        )
+
+        # Decode GPT-2 output
+        gpt2_response = tokenizer_gpt2.decode(outputs_gpt2[0], skip_special_tokens=True)
+
+        return f"SQL Query: {sql_query}\n\nResults:\n{gpt2_response}"
+    else:
+        return "No results found or error fetching data from Shopify."
+
+def main():
+    # Gradio interface setup
+    iface = gr.Interface(
+        fn=generate_sql,
+        inputs="text",
+        outputs="text",
+        title="Natural Language to SQL Query Generator",
+        description="Enter a natural language query to generate the corresponding SQL query and display results."
+    )
+    iface.launch()
+
+if __name__ == "__main__":
+    main()