jarif's picture
Update app.py
7c56e78 verified
raw
history blame
2.98 kB
import os
import logging
import streamlit as st
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain_huggingface import HuggingFacePipeline
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
# Configure logging
logging.basicConfig(level=logging.DEBUG)
def load_vector_store():
# Ensure the directory exists
persist_directory = "./chroma_db"
if not os.path.exists(persist_directory):
logging.error(f"The directory '{persist_directory}' does not exist. Please run the ingestion script.")
st.error(f"The directory '{persist_directory}' does not exist. Please run the ingestion script.")
return None
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vector_store = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
return vector_store
def load_llm():
checkpoint = "LaMini-T5-738M"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint)
pipe = pipeline(
'text2text-generation',
model=model,
tokenizer=tokenizer,
max_length=1024, # Increase max_length to allow longer output
do_sample=True,
temperature=0, # Adjust temperature to control randomness
top_p=0.9 # Adjust top_p to control diversity of the output
)
return HuggingFacePipeline(pipeline=pipe)
def process_answer(question):
try:
vector_store = load_vector_store()
if vector_store is None:
return "Vector store not found. Please run the ingestion script.", {}
llm = load_llm()
qa = RetrievalQA.from_chain_type(
llm=llm,
chain_type="stuff",
retriever=vector_store.as_retriever(),
return_source_documents=True
)
result = qa.invoke(question)
answer = result['result']
return answer, result
except Exception as e:
logging.error(f"An error occurred while processing the answer: {e}")
st.error(f"An error occurred while processing the answer: {e}")
return "An error occurred while processing your request.", {}
def main():
st.title("Search Your PDF πŸ“šπŸ“")
with st.expander("About the App"):
st.markdown(
"""
This is a Generative AI-powered Question and Answering app that responds to questions about your PDF File.
"""
)
question = st.text_area("Enter your Question")
if st.button("Ask"):
st.info("Your Question: " + question)
st.info("Your Answer")
try:
answer, metadata = process_answer(question)
st.write(answer)
st.write(metadata)
except Exception as e:
st.error(f"An unexpected error occurred: {e}")
if __name__ == '__main__':
main()