documentQABot / app.py
KrishnaKumar23's picture
Update app.py
958a1fd
raw
history blame
2.87 kB
import streamlit as st
from streamlit_lottie import st_lottie
import fitz # PyMuPDF
import requests
import os, shutil
import sidebar
import llm_model
@st.cache_data(experimental_allow_widgets=True)
def index_document(uploaded_file):
if uploaded_file is not None:
# Specify the folder path where you want to store the uploaded file in the 'assets' folder
assets_folder = "assets/uploaded_files"
if not os.path.exists(assets_folder):
os.makedirs(assets_folder)
# Save the uploaded file to the specified folder
file_path = os.path.join(assets_folder, uploaded_file.name)
with open(file_path, "wb") as f:
f.write(uploaded_file.getvalue())
file_name = os.path.join(assets_folder, uploaded_file.name)
st.success(f"File '{file_name}' uploaded !")
with st.spinner("Indexing document... This is a free CPU version and may take a while⏳"):
llm_model.create_vector_db(file_name, instructor_embeddings)
return file_name
else:
return None
def load_lottieurl(url: str):
r = requests.get(url)
if r.status_code != 200:
return None
return r.json()
def is_query_valid(query: str) -> bool:
if not query:
st.error("Please enter a question!")
return False
return True
# Function to load model parameters
@st.cache_resource()
def load_model():
return llm_model.load_model_params()
st.set_page_config(page_title="Document QA Bot")
lottie_book = load_lottieurl("https://assets4.lottiefiles.com/temp/lf20_aKAfIn.json")
st_lottie(lottie_book, speed=1, height=200, key="initial")
# Place the title below the Lottie animation
st.title("Document Q&A Bot 🤖")
# Left Sidebar
sidebar.sidebar()
# st.sidebar.header("Upload PDF")
# load model parameters
llm, instructor_embeddings = load_model()
# Upload file through Streamlit
uploaded_file = st.file_uploader("Upload a file", type=["pdf", "doc", "docx", "txt"])
filename = index_document(uploaded_file)
print(filename)
if not filename:
st.stop()
with st.form(key="qa_form"):
query = st.text_area("Ask a question about the document")
submit = st.form_submit_button("Submit")
if submit:
if not is_query_valid(query):
st.stop()
# Output Columns
answer_col, sources_col = st.columns(2)
qa_chain = llm_model.document_parser(instructor_embeddings, llm)
result = qa_chain(query)
with answer_col:
st.markdown("#### Answer")
st.markdown(result["result"])
with sources_col:
st.markdown("#### Sources")
if not ("i don't know" in result["result"].lower()):
for source in result["source_documents"]:
st.markdown(source.page_content)
st.markdown(source.metadata["source"])
st.markdown("--------------------------")