Upload sql_model.py

sql_model.py

@@ -0,0 +1,111 @@
+# from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+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)
+        self.tokenizer = AutoTokenizer.from_pretrained("google/gemma-7b-it")
+        self.model = AutoModelForCausalLM.from_pretrained("google/gemma-7b-it",)
+
+    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 :
+
+        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