Spaces:
				
			
			
	
			
			
		Sleeping
		
	
	
	
			
			
	
	
	
	
		
		
		Sleeping
		
	Create app.py
Browse files
    	
        app.py
    ADDED
    
    | @@ -0,0 +1,170 @@ | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | |
|  | 
|  | |
| 1 | 
            +
            import streamlit as st
         | 
| 2 | 
            +
            import os
         | 
| 3 | 
            +
            import zipfile
         | 
| 4 | 
            +
            import shutil
         | 
| 5 | 
            +
            from io import BytesIO
         | 
| 6 | 
            +
            from PyPDF2 import PdfReader
         | 
| 7 | 
            +
            from langchain.text_splitter import RecursiveCharacterTextSplitter
         | 
| 8 | 
            +
            from langchain_community.embeddings import HuggingFaceEmbeddings
         | 
| 9 | 
            +
            from langchain_community.vectorstores import FAISS
         | 
| 10 | 
            +
            from langchain_community.docstore.in_memory import InMemoryDocstore
         | 
| 11 | 
            +
            from langchain_community.llms import HuggingFaceHub
         | 
| 12 | 
            +
            from langchain.chains import RetrievalQA
         | 
| 13 | 
            +
            from langchain.prompts import PromptTemplate
         | 
| 14 | 
            +
            import faiss
         | 
| 15 | 
            +
            import uuid
         | 
| 16 | 
            +
            from dotenv import load_dotenv
         | 
| 17 | 
            +
             | 
| 18 | 
            +
            # Load environment variables
         | 
| 19 | 
            +
            load_dotenv()
         | 
| 20 | 
            +
            HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
         | 
| 21 | 
            +
            RAG_ACCESS_KEY = os.getenv("RAG_ACCESS_KEY")
         | 
| 22 | 
            +
             | 
| 23 | 
            +
            # Initialize session state
         | 
| 24 | 
            +
            if "vectorstore" not in st.session_state:
         | 
| 25 | 
            +
                st.session_state.vectorstore = None
         | 
| 26 | 
            +
            if "history" not in st.session_state:
         | 
| 27 | 
            +
                st.session_state.history = []
         | 
| 28 | 
            +
            if "authenticated" not in st.session_state:
         | 
| 29 | 
            +
                st.session_state.authenticated = False
         | 
| 30 | 
            +
             | 
| 31 | 
            +
            # Sidebar
         | 
| 32 | 
            +
            with st.sidebar:
         | 
| 33 | 
            +
                st.header("RAG Control Panel")
         | 
| 34 | 
            +
                api_key_input = st.text_input("Enter RAG Access Key", type="password")
         | 
| 35 | 
            +
                
         | 
| 36 | 
            +
                # Authentication
         | 
| 37 | 
            +
                if st.button("Authenticate"):
         | 
| 38 | 
            +
                    if api_key_input == RAG_ACCESS_KEY:
         | 
| 39 | 
            +
                        st.session_state.authenticated = True
         | 
| 40 | 
            +
                        st.success("Authentication successful!")
         | 
| 41 | 
            +
                    else:
         | 
| 42 | 
            +
                        st.error("Invalid API key.")
         | 
| 43 | 
            +
                
         | 
| 44 | 
            +
                # File uploader
         | 
| 45 | 
            +
                if st.session_state.authenticated:
         | 
| 46 | 
            +
                    input_type = st.selectbox("Select Input Type", ["Single PDF", "Folder/Zip of PDFs"])
         | 
| 47 | 
            +
                    input_data = None
         | 
| 48 | 
            +
                    if input_type == "Single PDF":
         | 
| 49 | 
            +
                        input_data = st.file_uploader("Upload a PDF file", type=["pdf"])
         | 
| 50 | 
            +
                    else:
         | 
| 51 | 
            +
                        input_data = st.file_uploader("Upload a folder or zip of PDFs", type=["zip"])
         | 
| 52 | 
            +
                    
         | 
| 53 | 
            +
                    if st.button("Process Files") and input_data is not None:
         | 
| 54 | 
            +
                        with st.spinner("Processing files..."):
         | 
| 55 | 
            +
                            vector_store = process_input(input_type, input_data)
         | 
| 56 | 
            +
                            st.session_state.vectorstore = vector_store
         | 
| 57 | 
            +
                            st.success("Files processed successfully. You can now ask questions.")
         | 
| 58 | 
            +
                
         | 
| 59 | 
            +
                # Display chat history
         | 
| 60 | 
            +
                st.subheader("Chat History")
         | 
| 61 | 
            +
                for i, (q, a) in enumerate(st.session_state.history):
         | 
| 62 | 
            +
                    st.write(f"**Q{i+1}:** {q}")
         | 
| 63 | 
            +
                    st.write(f"**A{i+1}:** {a}")
         | 
| 64 | 
            +
                    st.markdown("---")
         | 
| 65 | 
            +
             | 
| 66 | 
            +
            # Main app
         | 
| 67 | 
            +
            def main():
         | 
| 68 | 
            +
                st.title("RAG Q&A App with Mistral AI")
         | 
| 69 | 
            +
                
         | 
| 70 | 
            +
                if not st.session_state.authenticated:
         | 
| 71 | 
            +
                    st.warning("Please authenticate with your API key in the sidebar.")
         | 
| 72 | 
            +
                    return
         | 
| 73 | 
            +
                
         | 
| 74 | 
            +
                if st.session_state.vectorstore is None:
         | 
| 75 | 
            +
                    st.info("Please upload and process a PDF or folder/zip of PDFs in the sidebar.")
         | 
| 76 | 
            +
                    return
         | 
| 77 | 
            +
                
         | 
| 78 | 
            +
                query = st.text_input("Enter your question:")
         | 
| 79 | 
            +
                if st.button("Submit") and query:
         | 
| 80 | 
            +
                    with st.spinner("Generating answer..."):
         | 
| 81 | 
            +
                        answer = answer_question(st.session_state.vectorstore, query)
         | 
| 82 | 
            +
                        st.session_state.history.append((query, answer))
         | 
| 83 | 
            +
                        st.write("**Answer:**", answer)
         | 
| 84 | 
            +
             | 
| 85 | 
            +
            def process_input(input_type, input_data):
         | 
| 86 | 
            +
                # Create uploads directory
         | 
| 87 | 
            +
                os.makedirs("uploads", exist_ok=True)
         | 
| 88 | 
            +
                
         | 
| 89 | 
            +
                documents = ""
         | 
| 90 | 
            +
                if input_type == "Single PDF":
         | 
| 91 | 
            +
                    pdf_reader = PdfReader(input_data)
         | 
| 92 | 
            +
                    for page in pdf_reader.pages:
         | 
| 93 | 
            +
                        documents += page.extract_text() or ""
         | 
| 94 | 
            +
                else:
         | 
| 95 | 
            +
                    # Handle zip file
         | 
| 96 | 
            +
                    zip_path = "uploads/uploaded.zip"
         | 
| 97 | 
            +
                    with open(zip_path, "wb") as f:
         | 
| 98 | 
            +
                        f.write(input_data.getvalue())
         | 
| 99 | 
            +
                    with zipfile.ZipFile(zip_path, "r") as zip_ref:
         | 
| 100 | 
            +
                        zip_ref.extractall("uploads/extracted")
         | 
| 101 | 
            +
                    
         | 
| 102 | 
            +
                    # Process all PDFs in extracted folder
         | 
| 103 | 
            +
                    for root, _, files in os.walk("uploads/extracted"):
         | 
| 104 | 
            +
                        for file in files:
         | 
| 105 | 
            +
                            if file.endswith(".pdf"):
         | 
| 106 | 
            +
                                pdf_path = os.path.join(root, file)
         | 
| 107 | 
            +
                                pdf_reader = PdfReader(pdf_path)
         | 
| 108 | 
            +
                                for page in pdf_reader.pages:
         | 
| 109 | 
            +
                                    documents += page.extract_text() or ""
         | 
| 110 | 
            +
                    
         | 
| 111 | 
            +
                    # Clean up extracted files
         | 
| 112 | 
            +
                    shutil.rmtree("uploads/extracted", ignore_errors=True)
         | 
| 113 | 
            +
                    os.remove(zip_path)
         | 
| 114 | 
            +
                
         | 
| 115 | 
            +
                # Split text
         | 
| 116 | 
            +
                text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
         | 
| 117 | 
            +
                texts = text_splitter.split_text(documents)
         | 
| 118 | 
            +
                
         | 
| 119 | 
            +
                # Create embeddings
         | 
| 120 | 
            +
                hf_embeddings = HuggingFaceEmbeddings(
         | 
| 121 | 
            +
                    model_name="sentence-transformers/all-mpnet-base-v2",
         | 
| 122 | 
            +
                    model_kwargs={'device': 'cpu'}
         | 
| 123 | 
            +
                )
         | 
| 124 | 
            +
                
         | 
| 125 | 
            +
                # Initialize FAISS
         | 
| 126 | 
            +
                dimension = len(hf_embeddings.embed_query("sample text"))
         | 
| 127 | 
            +
                index = faiss.IndexFlatL2(dimension)
         | 
| 128 | 
            +
                vector_store = FAISS(
         | 
| 129 | 
            +
                    embedding_function=hf_embeddings,
         | 
| 130 | 
            +
                    index=index,
         | 
| 131 | 
            +
                    docstore=InMemoryDocstore({}),
         | 
| 132 | 
            +
                    index_to_docstore_id={}
         | 
| 133 | 
            +
                )
         | 
| 134 | 
            +
                
         | 
| 135 | 
            +
                # Add texts to vector store
         | 
| 136 | 
            +
                uuids = [str(uuid.uuid4()) for _ in range(len(texts))]
         | 
| 137 | 
            +
                vector_store.add_texts(texts, ids=uuids)
         | 
| 138 | 
            +
                
         | 
| 139 | 
            +
                # Save vector store locally
         | 
| 140 | 
            +
                vector_store.save_local("vectorstore/faiss_index")
         | 
| 141 | 
            +
                
         | 
| 142 | 
            +
                return vector_store
         | 
| 143 | 
            +
             | 
| 144 | 
            +
            def answer_question(vectorstore, query):
         | 
| 145 | 
            +
                llm = HuggingFaceHub(
         | 
| 146 | 
            +
                    repo_id="mistralai/Mistral-7B-Instruct-v0.1",
         | 
| 147 | 
            +
                    model_kwargs={"temperature": 0.7, "max_length": 512},
         | 
| 148 | 
            +
                    huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN
         | 
| 149 | 
            +
                )
         | 
| 150 | 
            +
                
         | 
| 151 | 
            +
                retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
         | 
| 152 | 
            +
                
         | 
| 153 | 
            +
                prompt_template = PromptTemplate(
         | 
| 154 | 
            +
                    template="Use the provided context to answer the question concisely:\n\nContext: {context}\n\nQuestion: {question}\n\nAnswer:",
         | 
| 155 | 
            +
                    input_variables=["context", "question"]
         | 
| 156 | 
            +
                )
         | 
| 157 | 
            +
                
         | 
| 158 | 
            +
                qa_chain = RetrievalQA.from_chain_type(
         | 
| 159 | 
            +
                    llm=llm,
         | 
| 160 | 
            +
                    chain_type="stuff",
         | 
| 161 | 
            +
                    retriever=retriever,
         | 
| 162 | 
            +
                    return_source_documents=False,
         | 
| 163 | 
            +
                    chain_type_kwargs={"prompt": prompt_template}
         | 
| 164 | 
            +
                )
         | 
| 165 | 
            +
                
         | 
| 166 | 
            +
                result = qa_chain({"query": query})
         | 
| 167 | 
            +
                return result["result"].split("Answer:")[-1].strip()
         | 
| 168 | 
            +
             | 
| 169 | 
            +
            if __name__ == "__main__":
         | 
| 170 | 
            +
                main()
         | 
