Update app.py

app.py

@@ -1,66 +1,66 @@
-import logging
-import gradio as gr
-from sql_generator import SQLGenerator
-from intent_classifier import IntentClassifier
-from rag_system import RAGSystem
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import requests
+from config import ACCESS_TOKEN, SHOP_NAME
 
-# Initialize logging
-logging.basicConfig(
-    level=logging.INFO,  # Set the logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
-    format='%(asctime)s - %(levelname)s - %(message)s'  # Format for the log messages
-)
-
-class UnifiedSystem:
-    def __init__(self): 
-        self.sql_generator = SQLGenerator()
-        self.intent_classifier = IntentClassifier()
-        self.rag_system = RAGSystem()
+class SQLGenerator:
+    def __init__(self):
+        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 process_query(self, query):
-        logging.info(f"Processing query: {query}")  # Log the incoming query
-        intent, confidence = self.intent_classifier.classify(query)
+    def generate_query(self, natural_language_query):
+        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)
+        );
+        """
         
-        logging.info(f"Classified intent: {intent} with confidence: {confidence:.2f}")
-
-        if intent == "database_query":
-            sql_query = self.sql_generator.generate_query(query)  # Assuming this method is correct
-            products = self.sql_generator.fetch_shopify_data("products")
-            
-            if products and 'products' in products:
-                results = "\n".join([
-                    f"Title: {p['title']}, Vendor: {p['vendor']}" 
-                    for p in products['products']
-                ])
-                return f"Intent: Database Query (Confidence: {confidence:.2f})\n\n" \
-                       f"SQL Query: {sql_query}\n\nResults:\n{results}"
-            else:
-                logging.warning("No results found or error fetching data from Shopify.")
-                return "No results found or error fetching data from Shopify."
-                
-        elif intent == "product_description":
-            rag_response = self.rag_system.process_query(query)
-            return f"Intent: Product Description (Confidence: {confidence:.2f})\n\n" \
-                   f"Response: {rag_response}"
+        prompt = f"""### Task: Generate a SQL query to answer the following question.
+        ### Database Schema:
+        {schema_info}
+        ### Question: {natural_language_query}
+        ### SQL Query:"""
         
-        logging.error("Intent not recognized.")
-        return "Intent not recognized."
-
-def create_interface():
-    system = UnifiedSystem()
-    
-    iface = gr.Interface(
-        fn=system.process_query,
-        inputs=gr.Textbox(
-            label="Enter your query",
-            placeholder="e.g., 'Show me all T-shirts' or 'Describe the product features'"
-        ),
-        outputs=gr.Textbox(label="Response"),
-        title="Unified Query Processing System",
-        description="Enter a natural language query to search products or get descriptions."
-    )
+        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
+        )
+        
+        return self.tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
     
-    return iface
-
-if __name__ == "__main__":
-    iface = create_interface()
-    iface.launch()
+    def fetch_shopify_data(self, endpoint):
+        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