Update app.py
Browse files
app.py
CHANGED
@@ -1,70 +1,81 @@
|
|
1 |
-
import os
|
2 |
-
import logging
|
3 |
-
import streamlit as st
|
4 |
-
from langchain.embeddings import HuggingFaceEmbeddings
|
5 |
-
from langchain.vectorstores import Chroma
|
6 |
-
from langchain.chains import RetrievalQA
|
7 |
-
from
|
8 |
-
|
9 |
-
|
10 |
-
logging
|
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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import os
|
2 |
+
import logging
|
3 |
+
import streamlit as st
|
4 |
+
from langchain.embeddings import HuggingFaceEmbeddings
|
5 |
+
from langchain.vectorstores import Chroma
|
6 |
+
from langchain.chains import RetrievalQA
|
7 |
+
from langchain_huggingface import HuggingFacePipeline
|
8 |
+
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
|
9 |
+
|
10 |
+
# Configure logging
|
11 |
+
logging.basicConfig(level=logging.DEBUG)
|
12 |
+
|
13 |
+
def load_vector_store():
|
14 |
+
# Ensure the directory exists
|
15 |
+
persist_directory = "./chroma_db"
|
16 |
+
if not os.path.exists(persist_directory):
|
17 |
+
logging.error(f"The directory '{persist_directory}' does not exist. Please run the ingestion script.")
|
18 |
+
st.error(f"The directory '{persist_directory}' does not exist. Please run the ingestion script.")
|
19 |
+
return None
|
20 |
+
|
21 |
+
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
|
22 |
+
vector_store = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
|
23 |
+
return vector_store
|
24 |
+
|
25 |
+
def load_llm():
|
26 |
+
checkpoint = "LaMini-T5-738M"
|
27 |
+
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
|
28 |
+
model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)
|
29 |
+
pipe = pipeline(
|
30 |
+
'text2text-generation',
|
31 |
+
model=model,
|
32 |
+
tokenizer=tokenizer,
|
33 |
+
max_length=256,
|
34 |
+
do_sample=True,
|
35 |
+
temperature=0.3,
|
36 |
+
top_p=0.95
|
37 |
+
)
|
38 |
+
return HuggingFacePipeline(pipeline=pipe)
|
39 |
+
|
40 |
+
def process_answer(question):
|
41 |
+
try:
|
42 |
+
vector_store = load_vector_store()
|
43 |
+
if vector_store is None:
|
44 |
+
return "Vector store not found. Please run the ingestion script.", {}
|
45 |
+
|
46 |
+
llm = load_llm()
|
47 |
+
qa = RetrievalQA.from_chain_type(
|
48 |
+
llm=llm,
|
49 |
+
chain_type="stuff",
|
50 |
+
retriever=vector_store.as_retriever(),
|
51 |
+
return_source_documents=True
|
52 |
+
)
|
53 |
+
result = qa.invoke(question)
|
54 |
+
answer = result['result']
|
55 |
+
return answer, result
|
56 |
+
except Exception as e:
|
57 |
+
logging.error(f"An error occurred while processing the answer: {e}")
|
58 |
+
st.error(f"An error occurred while processing the answer: {e}")
|
59 |
+
return "An error occurred while processing your request.", {}
|
60 |
+
|
61 |
+
def main():
|
62 |
+
st.title("Search Your PDF ππ")
|
63 |
+
with st.expander("About the App"):
|
64 |
+
st.markdown(
|
65 |
+
"""
|
66 |
+
This is a Generative AI powered Question and Answering app that responds to questions about your PDF File.
|
67 |
+
"""
|
68 |
+
)
|
69 |
+
question = st.text_area("Enter your Question")
|
70 |
+
if st.button("Ask"):
|
71 |
+
st.info("Your Question: " + question)
|
72 |
+
st.info("Your Answer")
|
73 |
+
try:
|
74 |
+
answer, metadata = process_answer(question)
|
75 |
+
st.write(answer)
|
76 |
+
st.write(metadata)
|
77 |
+
except Exception as e:
|
78 |
+
st.error(f"An unexpected error occurred: {e}")
|
79 |
+
|
80 |
+
if __name__ == '__main__':
|
81 |
+
main()
|