Create app.py

app.py

@@ -0,0 +1,202 @@
+import os
+import streamlit as st
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain.chains import RetrievalQA
+from langchain_community.vectorstores import FAISS
+from langchain_core.prompts import PromptTemplate
+from langchain_huggingface import HuggingFaceEndpoint
+from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+
+# Load environment variables
+from dotenv import load_dotenv, find_dotenv
+load_dotenv(find_dotenv())
+
+# Constants
+DATA_PATH = "data/"
+DB_FAISS_PATH = "vectorstore/db_faiss"
+HUGGINGFACE_REPO_ID = "mistralai/Mistral-7B-Instruct-v0.3"
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+# Custom prompt template
+CUSTOM_PROMPT_TEMPLATE = """
+Use the pieces of information provided in the context to answer user's question.
+If you dont know the answer, just say that you dont know, dont try to make up an answer.
+
+Dont provide anything out of the given context
+
+Context: {context}
+Question: {question}
+
+Start the answer directly. No small talk please.
+"""
+
+def load_pdf_files(data_path):
+    try:
+        loader = DirectoryLoader(data_path,
+                                glob='*.pdf',
+                                loader_cls=PyPDFLoader)
+        documents = loader.load()
+        return documents
+    except Exception as e:
+        st.error(f"Error loading PDF files: {e}")
+        return []
+
+def create_chunks(extracted_data):
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500,
+                                                chunk_overlap=50)
+    text_chunks = text_splitter.split_documents(extracted_data)
+    return text_chunks
+
+def get_embedding_model():
+    embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    return embedding_model
+
+def create_vectorstore():
+    if not os.path.exists(DATA_PATH):
+        os.makedirs(DATA_PATH)
+        st.warning(f"Created empty data directory at {DATA_PATH}. Please upload PDF files.")
+        return None
+    
+    documents = load_pdf_files(data=DATA_PATH)
+    if not documents:
+        st.warning("No PDF files found in data directory. Please upload some PDFs.")
+        return None
+    
+    st.info(f"Loaded {len(documents)} PDF pages")
+    text_chunks = create_chunks(extracted_data=documents)
+    st.info(f"Created {len(text_chunks)} text chunks")
+    
+    embedding_model = get_embedding_model()
+    
+    if not os.path.exists(os.path.dirname(DB_FAISS_PATH)):
+        os.makedirs(os.path.dirname(DB_FAISS_PATH))
+    
+    db = FAISS.from_documents(text_chunks, embedding_model)
+    db.save_local(DB_FAISS_PATH)
+    st.success(f"Created vector store at {DB_FAISS_PATH}")
+    return db
+
+@st.cache_resource
+def get_vectorstore():
+    if os.path.exists(DB_FAISS_PATH):
+        embedding_model = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
+        try:
+            db = FAISS.load_local(DB_FAISS_PATH, embedding_model, allow_dangerous_deserialization=True)
+            return db
+        except Exception as e:
+            st.error(f"Error loading vector store: {e}")
+            return None
+    else:
+        st.warning("Vector store not found. Please create it first.")
+        return None
+
+def set_custom_prompt():
+    prompt = PromptTemplate(template=CUSTOM_PROMPT_TEMPLATE, input_variables=["context", "question"])
+    return prompt
+
+def load_llm():
+    if not HF_TOKEN:
+        st.error("HF_TOKEN not found. Please set it in your environment variables.")
+        return None
+    
+    try:
+        llm = HuggingFaceEndpoint(
+            repo_id=HUGGINGFACE_REPO_ID,
+            task="text-generation",
+            temperature=0.5,
+            model_kwargs={
+                "token": HF_TOKEN,
+                "max_length": 512
+            }
+        )
+        return llm
+    except Exception as e:
+        st.error(f"Error loading LLM: {e}")
+        return None
+
+def upload_pdf():
+    uploaded_files = st.file_uploader("Upload PDF files", type="pdf", accept_multiple_files=True)
+    if uploaded_files:
+        for uploaded_file in uploaded_files:
+            with open(os.path.join(DATA_PATH, uploaded_file.name), "wb") as f:
+                f.write(uploaded_file.getbuffer())
+        st.success(f"Uploaded {len(uploaded_files)} files to {DATA_PATH}")
+        return True
+    return False
+
+def main():
+    st.title("PDF Question Answering System")
+    
+    # Sidebar
+    st.sidebar.title("Settings")
+    page = st.sidebar.radio("Choose an action", ["Upload PDFs", "Create Vector Store", "Chat with Documents"])
+    
+    if page == "Upload PDFs":
+        st.header("Upload PDF Files")
+        st.info("Upload PDF files that will be used for question answering")
+        if upload_pdf():
+            st.info("Now go to 'Create Vector Store' to process your documents")
+    
+    elif page == "Create Vector Store":
+        st.header("Create Vector Store")
+        st.info("This will process your PDF files and create embeddings")
+        if st.button("Create Vector Store"):
+            with st.spinner("Processing documents..."):
+                create_vectorstore()
+    
+    elif page == "Chat with Documents":
+        st.header("Ask Questions About Your Documents")
+        
+        if 'messages' not in st.session_state:
+            st.session_state.messages = []
+        
+        for message in st.session_state.messages:
+            st.chat_message(message['role']).markdown(message['content'])
+        
+        prompt = st.chat_input("Ask a question about your documents")
+        
+        if prompt:
+            st.chat_message('user').markdown(prompt)
+            st.session_state.messages.append({'role': 'user', 'content': prompt})
+            
+            vectorstore = get_vectorstore()
+            if vectorstore is None:
+                st.error("Vector store not available. Please create it first.")
+                return
+            
+            llm = load_llm()
+            if llm is None:
+                return
+            
+            try:
+                with st.spinner("Thinking..."):
+                    qa_chain = RetrievalQA.from_chain_type(
+                        llm=llm,
+                        chain_type="stuff",
+                        retriever=vectorstore.as_retriever(search_kwargs={'k': 3}),
+                        return_source_documents=True,
+                        chain_type_kwargs={'prompt': set_custom_prompt()}
+                    )
+                    
+                    response = qa_chain.invoke({'query': prompt})
+                    
+                    result = response["result"]
+                    source_documents = response["source_documents"]
+                    
+                    # Format source documents more cleanly
+                    source_docs_text = "\n\n**Source Documents:**\n"
+                    for i, doc in enumerate(source_documents, 1):
+                        source_docs_text += f"{i}. Page {doc.metadata.get('page', 'N/A')}: {doc.page_content[:200]}...\n\n"
+                    
+                    result_to_show = f"{result}\n{source_docs_text}"
+                    
+                    st.chat_message('assistant').markdown(result_to_show)
+                    st.session_state.messages.append({'role': 'assistant', 'content': result_to_show})
+                    
+            except Exception as e:
+                st.error(f"Error: {str(e)}")
+                st.error("Please check your HuggingFace token and model access permissions")
+
+if __name__ == "__main__":
+    main()