Upload 4 files

nlp/sql_model.py

@@ -1,113 +1,105 @@
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-# import os
-# from transformers import AutoTokenizer, AutoModelForCausalLM
+import re
+from huggingface_hub import InferenceClient
 
 class SQLModel:
-    def __init__(self, model_name="google/flan-t5-base"):
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-        # hf_token = os.environ.get("HF_HUB_TOKEN")
-        # self.tokenizer = AutoTokenizer.from_pretrained("google/gemma-2b-it",use_auth_token=hf_token)
-        # self.model = AutoModelForCausalLM.from_pretrained("google/gemma-2b-it",use_auth_token=hf_token)
+    def __init__(self, model_name="HuggingFaceH4/zephyr-7b-beta"):
+        self.client = InferenceClient(model_name)
 
     def generate_sql(self, natural_language_query):
-        input_text = f"""You are a highly skilled SQL translator. Your task is to convert natural language descriptions of data queries into correct and optimized SQL statements.
-        Here is the schema information for our database :
+        prompt = (
+            "You are a highly skilled SQL translator. Your task is to convert natural language descriptions of data queries "
+            "into correct and optimized SQL statements.\n\n"
+            "Here is the schema information for our database :\n\n"
+            "Table: Employees\n"
+            "- id (INT)\n"
+            "- NAME (VARCHAR)\n"
+            "- Department (VARCHAR)\n"
+            "- Salary (INT)\n"
+            "- Hire_Date (DATE)\n\n"
+            "Table: Departments\n"
+            "- ID (INT)\n"
+            "- Name (VARCHAR)\n"
+            "- Manager (VARCHAR)\n\n"
+            "Here are a few examples:\n\n"
+            "1. **Input**: \"Show me all employees in the Sales department.\"\n"
+            "**Output**:\n\n"
+            "        SELECT *\n"
+            "        FROM Employees\n"
+            "        WHERE Department = 'Sales';\n\n"
+            "2. **Input**: \"Who is the manager of the Engineering department?\"\n"
+            "**Output**:\n\n"
+            "        SELECT Manager\n"
+            "        FROM Departments\n"
+            "        WHERE Name = 'Engineering';\n\n"
+            "3. **Input**: \"List all employees hired after 2021-01-01.\"\n"
+            "**Output**:\n\n"
+            "        SELECT *\n"
+            "        FROM Employees\n"
+            "        WHERE Hire_Date > '2021-01-01';\n\n"
+            "4. **Input**: \"What is the total salary expense for the Marketing department?\"\n"
+            "**Output**:\n\n"
+            "        SELECT SUM(Salary)\n"
+            "        FROM Employees\n"
+            "        WHERE Department = 'Marketing';\n\n"
+            "5. **Input**: \"Find the average salary of employees in each department.\"\n"
+            "**Output**:\n\n"
+            "        SELECT Department, AVG(Salary) AS average_salary\n"
+            "        FROM Employees\n"
+            "        GROUP BY Department;\n\n"
+            "Please do not return additional text besides query.\n"
+            "Please only answer queries which makes sense for the given schema. Else just return - \"No information found\""
+        )
+
+        messages = [
+            {"role": "system", "content": prompt},
+            {"role": "user", "content": natural_language_query}
+        ]
+
+        result = self.client.chat_completion(
+            messages,
+            max_tokens=150,
+            stream=False,
+            temperature=0.7,
+            top_p=0.95,
+        )
+
+        # Initialize a variable to hold the extracted SQL text.
+        sql_query = ""
+
+        # Check if the result is a plain string.
+        if isinstance(result, str):
+            sql_query = result
+        # If the result is a list, iterate over its tokens.
+        elif isinstance(result, list):
+            for token in result:
+                if isinstance(token, str):
+                    sql_query += token
+                elif hasattr(token, "choices"):
+                    # Extract from the structured object.
+                    sql_query += token.choices[0].delta.content
+                else:
+                    sql_query += str(token)
+        # Otherwise, if it's an object with choices, extract its content.
+        elif hasattr(result, "choices"):
+            sql_query = result.choices[0].message.content
+        else:
+            sql_query = str(result)
+
+        # Optional: If the model output is in a markdown code block, extract only that content.
+        match = re.search(r"```sql(.*?)```", sql_query, re.DOTALL | re.IGNORECASE)
+        if match:
+            sql_query = match.group(1).strip()
+
+        # Remove both literal "\n" substrings and actual newline characters.
+        sql_query = sql_query.replace("\\n", " ").replace("\n", " ")
+        # Remove extra spaces.
+        sql_query = " ".join(sql_query.split())
+
+        # Extract only the SQL command: starting from the first occurrence of "select" to the first semicolon.
+        extraction_pattern = r"(?i)(select\s.*?;)"
+        extraction_match = re.search(extraction_pattern, sql_query, re.DOTALL)
+        if extraction_match:
+            sql_query = " ".join(extraction_match.group(1).split())
 
-        Table: Employees
-        - id (INT)
-        - NAME (VARCHAR)
-        - Department (VARCHAR)
-        - Salary (INT)
-        - Hire_Date (DATE)
-
-        Table: Departments
-        - ID (INT)
-        - Name (VARCHAR)
-        - Manager (VARCHAR)
-
-        Here are a few examples:
-
-        1. **Input**: "Show me all employees in the Sales department."
-        **Output**:
-
-                SELECT *
-                FROM Employees
-                WHERE Department = 'Sales';
-
-        2. **Input**: "Who is the manager of the Engineering department?"
-        **Output**:
-
-                SELECT Manager
-                FROM Departments
-                WHERE Name = 'Engineering';
-
-
-        3. **Input**: "List all employees hired after 2021-01-01."
-        **Output**:
-
-                SELECT *
-                FROM Employees
-                WHERE Hire_Date > '2021-01-01';
-
-
-        4. **Input**: "What is the total salary expense for the Marketing department?"
-        **Output**:
-
-                SELECT SUM(Salary)
-                FROM Employees
-                WHERE Department = 'Marketing';
-
-
-        5. **Input**: "Find the average salary of employees in each department."
-        **Output**:
-
-                SELECT Department, AVG(Salary) AS average_salary
-                FROM Employees
-                GROUP BY Department;
-
-        Please do not return additional text besides query.
-
-        Please only answer queries which makes sense for the given schema. Else just return - "No information found"
-
-        Now, translate the following natural language query into an syntactically correct SQL query:
-        **Input**: {natural_language_query}
-        **Output**:
-
-"""
-#         input_text = f"""
-# Translate the following natural language query into a syntactically correct SQL query using the provided database schema. Output only the SQL query with no additional text or explanation.
-
-# Database Schema:
-
-# Table: Employees
-# - id (INT)
-# - NAME (VARCHAR)
-# - Department (VARCHAR)
-# - Salary (INT)
-# - Hire_Date (DATE)
-
-# Table: Departments
-# - ID (INT)
-# - Name (VARCHAR)
-# - Manager (VARCHAR)
-
-# Examples:
-# 1. Natural Language Query: "List all employees who were hired after '2020-01-01'."
-#    Output: SELECT * FROM Employees WHERE Hire_Date > '2020-01-01';
-
-# 2. Natural Language Query: "Retrieve the names and salaries of employees in the 'Sales' department."
-#    Output: SELECT NAME, Salary FROM Employees WHERE Department = 'Sales';
-
-# Now, translate the following query:
-# {natural_language_query}
-# """
-        # input_text = f"translate English to SQL: {natural_language_query}"
-        # inputs = self.tokenizer(input_text, return_tensors="pt").input_ids
-        # outputs = self.model.generate(inputs, max_new_tokens=100, do_sample=False)
-        # sql_query = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        inputs = self.tokenizer(input_text, return_tensors="pt")
-        outputs = self.model.generate(**inputs)
-        sql_query = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
         print(sql_query)
         return sql_query