Created db.py and nlp.py for executing queries and converting queries to sql respectively

app/db.py

@@ -1,52 +1,18 @@
 import sqlite3
-
-DB_PATH = "data/database.sqlite"
-
-def get_connection():
-    return sqlite3.connect(DB_PATH)
-
-def create_tables():
-    conn = get_connection()
-    cursor = conn.cursor()
-
-    cursor.execute('''CREATE TABLE IF NOT EXISTS Employees(
-                    ID INTEGER PRIMARY KEY AUTOINCREMENT,
-                    Name TEXT NOT NULL,
-                    Department TEXT NOT NULL,
-                    Salary INTEGER NOT NULL,
-                    Hire_Date TEXT NOT NULL
-                    )
-    ''')
-
-    cursor.execute('''CREATE TABLE IF NOT EXISTS Departments(
-                ID INTEGER PRIMARY KEY AUTOINCREMENT,
-                Name TEXT NOT NULL,
-                Manager TEXT NOT NULL  
-                ) 
-    ''')
-
-    employees = [
-        (1, 'Alice', 'Sales', 50000, '2021-01-15'),
-        (2, 'Bob', 'Engineering', 70000, '2020-06-10'),
-        (3, 'Charlie', 'Marketing', 60000, '2022-03-20')
-    ]
-
-    departments = [
-        (1, 'Sales', 'Alice'),
-        (2, 'Engineering', 'Bob'),
-        (3, 'Marketing', 'Charlie')
-    ]
-
-    cursor.executemany('INSERT INTO Employees VALUES (?,?,?,?,?)', employees)
-    cursor.executemany('INSERT INTO Departments VALUES (?,?,?)', departments)
-
-    conn.commit()
-    conn.close()
-
-
-
-if __name__ == "__main__":
-    create_tables()
-
-
-
+from sqlite3 import Error
+
+class Database:
+    def __init__(self, db_path='data/database.sqlite'):
+        self.db_path = db_path
+
+    def execute_query(self, query):
+        try:
+            conn = sqlite3.connect(self.db_path)
+            cursor = conn.cursor()
+            cursor.execute(query)
+            results = cursor.fetchall()
+            columns = [desc[0] for desc in cursor.description]
+            conn.close()
+            return {"columns": columns, "data": results}
+        except Error as e:
+            return {"error": str(e)}

app/nlp.py

@@ -0,0 +1,14 @@
+from transformers import pipeline
+
+class NLPToSQL:
+    def __init__(self):
+        self.model = pipeline(
+            "text2text-generation",
+            model="mrm8488/t5-base-finetuned-wikiSQL",
+            tokenizer="t5-base"
+        )
+
+    def query_to_sql(self, user_query):
+        prompt = f"translate English to SQL: {user_query}"
+        result = self.model(prompt, max_length=200)
+        return result[0]['generated_text']