feat: add main menu options and scripts

app.py

@@ -5,6 +5,7 @@ import sys
 sys.path.insert(0, os.path.abspath(os.path.dirname(__file__)))
 
 import re
+import sqlite3
 import numpy as np
 import pandas as pd
 import streamlit as st
@@ -12,6 +13,7 @@ import requests
 
 from googletrans import Translator
 from langdetect import detect
+from sql_formatter.core import format_sql
 
 translator = Translator()
 
@@ -27,6 +29,28 @@ TEXT_2_SQL_API = os.environ.get(
 )
 
 
+@st.cache_resource
+def load_database():
+    db_conn = sqlite3.connect("resources/ai_app.db")
+    with open("resources/schema.sql", "r") as f:
+        db_conn.executescript(f.read())
+
+    return db_conn
+
+
+db_conn = load_database()
+
+
+def execute_sql(sql_query):
+    try:
+        cursor = db_conn.cursor()
+        cursor.execute(sql_query)
+        st.success("SQL query executed successfully!")
+        return cursor.fetchall()
+    except Exception as e:
+        st.info("Database is not supported")
+
+
 @st.cache_data
 def ask_text2sql(question, context):
     if detect(question) != "en":
@@ -62,11 +86,11 @@ def introduction():
         "🔍 This app allows you to explore the ability of Text to SQL model. The model is CodeLlama-13b finetuned using QLoRA on NSText2SQL dataset."
     )
     st.write(
-        "📈 The NSText2SQL dataset contains more than 290.000 training samples. Then, the model is evaluated on Spider and VMLP datasets."
+        "📈 The NSText2SQL dataset contains more than 290.000 training samples. Then, the model is evaluated on Spider and vMLP datasets."
     )
     st.write("📑 The other pages in this app include:")
     st.write(
-        "  - 📊 EDA Page: This page includes several visualizations to help you understand the two dataset: Spider and VMLP."
+        "  - 📊 EDA Page: This page includes several visualizations to help you understand the two dataset: Spider and vMLP."
     )
     st.write(
         "  - 💰 Text2SQL Page: This page allows you to generate SQL query from a given question and context."
@@ -159,7 +183,10 @@ def examples():
             if example_btns[idx]:
                 st.markdown("##### SQL query:")
                 query = ask_text2sql(row["question"], row["context"])
-                st.code(query, language="sql")
+                st.code(format_sql(query), language="sql")
+                result = execute_sql(query)
+                st.write("The result of the SQL query is:")
+                st.dataframe(pd.DataFrame(result))
 
 
 # Define a function for the Stock Prediction page
@@ -174,8 +201,13 @@ def interactive_demo():
     question_placeholder = st.empty()
     context = context_placeholder.text_area(
         "##### Context",
-        "CREATE TABLE head (name VARCHAR, born_state VARCHAR, age VARCHAR)",
+        """CREATE TABLE customer (id number, name text, gender text, age number, district_id number;
+CREATE TABLE registration (customer_id number, product_id number); 
+CREATE TABLE district (id number, name text, prefix text, province_id number); 
+CREATE TABLE province (id number, name text, code text)
+CREATE TABLE product (id number, category text, name text, description text, price number, duration number, data_amount number, voice_amount number, sms_amount number);""",
         key="context",
+        height=150,
     )
     question = question_placeholder.text_input(
         "##### Question",
@@ -189,7 +221,10 @@ def interactive_demo():
         query = ask_text2sql(question, context)
         st.write("The SQL query generated by the model is:")
         # Display the SQL query in a code block
-        st.code(query, language="sql")
+        st.code(format_sql(query), language="sql")
+        result = execute_sql(query)
+        st.write("The result of the SQL query is:")
+        st.dataframe(pd.DataFrame(result))
 
 
 # Define a function for the About page

requirements.txt

@@ -5,3 +5,4 @@ pandas
 googletrans==3.1.0a0
 langdetect
 requests
+sql-formatter