File size: 2,546 Bytes
c4fbb4c |
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 |
import streamlit as st
from langchain import HuggingFaceHub
from langchain.prompts import ChatPromptTemplate
from langchain_google_genai import ChatGoogleGenerativeAI, HarmBlockThreshold, HarmCategory
from langchain.chains.question_answering import load_qa_chain
from dotenv import load_dotenv
import os
import json
import pandas as pd
import numpy as np
import sqlite3
load_dotenv()
conn = sqlite3.connect('your_database.db')
c = conn.cursor()
def get_database_schema():
c.execute("SELECT name FROM sqlite_master WHERE type='table';")
tables = c.fetchall()
schema_dict = {}
for table in tables:
c.execute(f"PRAGMA table_info({table[0]})")
schema = c.fetchall()
schema_dict[table[0]] = schema
return schema_dict
def get_chain():
template = """
You are a SQLite query generator.
return output executable query only or empyt text .
output should be json with key "query" and tabel name with key "table_name"
Based on user question, generate SQLite Query for given database Schema.
Schema : \n{context}\n
Question : {question}
"""
hf_model = HuggingFaceHub(repo_id="Mistralai/Mistralaichat", model_kwargs={"temperature": 0.4})
prompt=ChatPromptTemplate.from_template(template)
chain=LLMChain(llm=hf_model.load_chat(),prompt=prompt)
return chain
def get_query(txt1):
start_index = txt1.find('{')
end_index = txt1.find('}') + 1
json_string = txt1[start_index:end_index]
data = json.loads(json_string)
return data
def query_with_database(query, tablename=None):
if tablename:
c.execute(query)
else:
c.execute(query)
query_output = c.fetchall()
df = pd.DataFrame(query_output, columns=[desc[1] for desc in c.description])
st.table(df)
def user_input(user_question):
schema = get_database_schema()
chain = get_chain()
response = chain.run({"context": schema, "question": user_question})
generated_query = get_query(response)
st.write("Generated Query : ", generated_query["query"])
query_with_database(generated_query["query"])
def main():
st.header("Chat With SQLite DATABASE - HUGGING FACE MISTRALAI/MISTRALAICHAT")
with st.sidebar:
st.write("Develop by : Rehan Shekh")
st.write("Model : Hugging Face Mistralai/Mistralaichat")
buttoclick=st.text_input(label="Ask from Database")
st.write("Demo Database only have one table :- User_details")
if buttoclick:
user_input(str(buttoclick))
if __name__=="__main__":
main() |