Update app.py

app.py

@@ -1,56 +1,47 @@
 import streamlit as st
 from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import PyPDFLoader
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
-from langchain.chains import ConversationalRetrievalChain
-from langchain_openai import ChatOpenAI
+from langchain.chains import RetrievalQA
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import tempfile
 import os
 
-st.set_page_config(page_title="Chat with Text", layout="wide")
+st.set_page_config(page_title="Document QA Bot")
 
-if "messages" not in st.session_state:
-    st.session_state.messages = []
 if "vector_store" not in st.session_state:
     st.session_state.vector_store = None
 
-def create_vector_store(text):
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    chunks = text_splitter.create_documents([text])
-    embeddings = HuggingFaceEmbeddings()
+def process_text(text):
+    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    chunks = splitter.create_documents([text])
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
     return FAISS.from_documents(chunks, embeddings)
 
-def get_qa_chain(vector_store):
-    llm = ChatOpenAI(temperature=0)
-    return ConversationalRetrievalChain.from_llm(
-        llm=llm,
-        retriever=vector_store.as_retriever(),
-        return_source_documents=True
-    )
-
-st.title("💬 Chat with Your Text")
-
-with st.sidebar:
-    api_key = st.text_input("OpenAI API Key", type="password")
-    if api_key:
-        os.environ["OPENAI_API_KEY"] = api_key
-    
-    text_input = st.text_area("Your Text Here", height=300)
-    if st.button("Process Text") and text_input and api_key:
-        with st.spinner("Processing text..."):
-            st.session_state.vector_store = create_vector_store(text_input)
-            st.success("Ready to chat!")
-
-if st.session_state.vector_store and api_key:
-    if question := st.chat_input("Ask your question about the text"):
-        chain = get_qa_chain(st.session_state.vector_store)
-        response = chain({"question": question, "chat_history": st.session_state.messages})
-        st.session_state.messages.append((question, response["answer"]))
-
-    for message in st.session_state.messages:
-        st.chat_message("user").write(message[0])
-        st.chat_message("assistant").write(message[1])
-
-elif not api_key:
-    st.warning("⚠️ Please enter your OpenAI API key")
-elif not st.session_state.vector_store:
-    st.info("👈 Please input your text and process it first")
+def process_pdf(file):
+    with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf') as tmp_file:
+        tmp_file.write(file.getvalue())
+        loader = PyPDFLoader(tmp_file.name)
+        pages = loader.load()
+        os.unlink(tmp_file.name)
+    return process_text('\n'.join(page.page_content for page in pages))
+
+st.title("Document QA Bot")
+
+uploaded_file = st.file_uploader("Upload Document", type=["txt", "pdf"])
+if uploaded_file:
+    with st.spinner("Processing document..."):
+        if uploaded_file.type == "text/plain":
+            text = uploaded_file.getvalue().decode()
+            st.session_state.vector_store = process_text(text)
+        else:
+            st.session_state.vector_store = process_pdf(uploaded_file)
+        st.success("Document processed!")
+
+if st.session_state.vector_store:
+    if question := st.chat_input("Ask a question about the document:"):
+        results = st.session_state.vector_store.similarity_search(question)
+        context = "\n".join(doc.page_content for doc in results)
+        st.chat_message("user").write(question)
+        st.chat_message("assistant").write(context)