File size: 3,661 Bytes
29e0766 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 |
from dotenv import find_dotenv, load_dotenv
load_dotenv(find_dotenv())
import os
import pandas as pd
import streamlit as st
from langchain_core.prompts import PromptTemplate
from pandasql import sqldf
from groq import Groq
template = """You are a powerful text-to-SQL model. Your job is to answer questions about a database. You are given a question and context regarding one or more tables. Dont add \n characters.
Do not include "SELECT short\_name, long\_name" this type of queries which have backslash in them.
You must output the SQL query that answers the question in a single line.
### Input:
`{question}`
### Context:
`{context}`
### Response:
"""
prompt = PromptTemplate.from_template(template=template)
client = Groq(
api_key=os.getenv("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"),
)
def groq_infer(prompt):
chat_completion = client.chat.completions.create(
messages=[
{
"role": "user",
"content": prompt,
}
],
model="mixtral-8x7b-32768",
)
print(chat_completion.choices[0].message.content)
return chat_completion.choices[0].message.content
# 1. Create cache_resource - To load the model
# infer - pipeline -> pipe()
def main():
st.set_page_config(page_title="Ask anything about database", page_icon="📊", layout="wide")
st.title("SQL Engineer")
col1, col2 = st.columns([2, 3])
with col1:
uploaded_file = st.file_uploader("Upload a CSV file", type="csv")
if uploaded_file is not None:
df = pd.read_csv(uploaded_file, encoding="latin1")
df.columns = df.columns.str.replace(r"[^a-zA-Z0-9_]", "", regex=True)
st.write("Here's a preview of your uploaded file:")
st.dataframe(df)
context = pd.io.sql.get_schema(df.reset_index(), "df").replace('"', "")
st.write("SQL Schema:")
st.code(context)
with col2:
if uploaded_file is not None:
question = st.text_input("Write a question about the data", key="question")
if st.button("Get Answer", key="get_answer"):
if question:
attempt = 0
max_attempts = 5
while attempt < max_attempts:
try:
input = {"context": context, "question": question}
formatted_prompt = prompt.invoke(input=input).text
response = groq_infer(formatted_prompt)
final = response.replace("`", "").replace("sql", "").strip()
st.text("Query performed")
st.code(final)
result = sqldf(final, locals())
st.write("Answer:")
st.dataframe(result)
break
except Exception as e:
attempt += 1
st.error(
f"Attempt {attempt}/{max_attempts} failed. Retrying..."
)
if attempt == max_attempts:
st.error(
"Unable to get the correct query, refresh app or try again later."
)
continue
else:
st.warning("Please enter a question before clicking 'Get Answer'.")
if __name__ == "__main__":
main() |