SqlAI / app.py
darkalchemist007's picture
Update app.py
152f485
import streamlit as st
from transformers import pipeline
import sqlite3
import numpy as np
import pandas as pd
#initialize chat history
if "messages" not in st.session_state:
st.session_state.messages=[]
## Function to read SQL data
def read_sql_query(sql, db):
st.write("in FUnction")
conn = sqlite3.connect(db)
cur = conn.cursor()
cur.execute(sql)
rows = cur.fetchall()
#for row in rows:
#st.write(row)
conn.close()
output_df = pd.DataFrame(rows)
st.table(output_df)
#return output_df
i_pwd= st.sidebar.text_input("Enter password", type="password")
if i_pwd=='password':
translator = pipeline(model= "saikiranmaddukuri/chat_to_sql0.17", max_new_tokens= 100)
df= pd.read_csv("Car_sales.csv")
st.title("Conversational AI: Text-2-SQL")
column_info = df.dtypes.to_dict()
column_df = pd.DataFrame(column_info.items(), columns=['Column Name', 'Data Type'])
st.sidebar.table(column_df)
create_table_statement= "CREATE TABLE car_sales (Manufacturer text, Model text, Sales_in_thousands float, __year_resale_value float, Vehicle_type text, Price_in_thousands float, Engine_size float, Horsepower float, Wheelbase float, Width float, Length float, Curb_weight float, Fuel_capacity float, Fuel_efficiency float, Latest_Launch text, Power_perf_factor float)"
# display chat messages from history
for message in st.session_state.messages:
with st.chat_message(message['role']):
st.markdown(message['content'])
# Step 1: User input
#with st.chat_message(name= 'assistant'):
#st.markdown('Ask question and press Enter')
#st.session_state.messages.append({"role":"assistant", "content": 'Ask question and press Enter' })
user_query= st.chat_input('Ask question and press ENTER.')
if user_query:
with st.chat_message('user'):
st.markdown(user_query)
st.session_state.messages.append({"role":"user", "content": user_query })
#user_query = st.text_input("")
# Step 2: Backend processing on Enter
#if user_query:
text='question:{}. table:{}'.format(user_query,create_table_statement.replace("text", 'VARCHAR') )
with st.spinner('Generating SQL. Please wait...'):
generated_sql= translator([text])
generated_sql= generated_sql[0]['generated_text']
i_gen_query = st.text_area("AI Generated SQL (edit if required)", generated_sql)
with st.chat_message('assistant'):
st.markdown(i_gen_query)
st.session_state.messages.append({"role":"assistant", "content": i_gen_query })
# Step 4: Process edited text on Enter
if st.button("Show results") and i_gen_query:
st.write("Button worked!")
st.markdown(i_gen_query)
#st.session_state.messages.append({"role":"assistant", "content": i_gen_query })
read_sql_query(i_gen_query, "Cars_db.sqlite")
# Step 5: Display final output
#st.write("Final Output:", final_output)
else:
st.warning("Enter correct password!")