Update app.py

app.py

@@ -11,7 +11,9 @@ from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.memory import ConversationBufferMemory
 from langchain.chains import ConversationalRetrievalChain
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-
+from docarray import Document
+from docarray import DocumentArray
+from sentence_transformers import SentenceTransformer
 
 # StreamHandler to intercept streaming output from the LLM.
 # This makes it appear that the Language Model is "typing"
@@ -39,36 +41,36 @@ def get_page_urls(url):
         return set()
 
 
+
 def get_url_content(url):
     response = requests.get(url)
     if url.endswith('.pdf'):
         pdf = io.BytesIO(response.content)
-        file = open('pdf.pdf', 'wb')
-        file.write(pdf.read())
-        file.close()
-        doc = fitz.open('pdf.pdf')
-        return (url, ''.join([text for page in doc for text in page.get_text("text")]))
+        doc = fitz.open(stream=pdf, filetype="pdf")
+        text = ''.join([page.get_text("text") for page in doc])
+        return Document(text=text, tags={'url': url})
     else:
         soup = BeautifulSoup(response.content, 'html.parser')
         content = soup.find_all('div', class_='wpb_content_element')
         text = ' '.join([c.get_text().strip() for c in content if c.get_text().strip() != ''])
-        return (url, text)
+        return Document(text=text, tags={'url': url})
+
+
 
 
 
 @st.cache_resource
 def get_retriever(urls):
-    all_content = [get_url_content(url) for url in urls]
-    documents = [Document(text=doc, metadata={'url': url}) for (url, doc) in all_content]
-
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=200)
-    docs = text_splitter.split_documents(documents)
+    documents = DocumentArray([get_url_content(url) for url in urls])
 
-    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+    # Load the model and encode document texts directly
+    model = SentenceTransformer('all-MiniLM-L6-v2')
+    documents.embeddings = model.encode(documents.texts, show_progress_bar=True)
 
-    db = DocArrayInMemorySearch.from_documents(docs, embeddings)
-    retriever = db.as_retriever(search_type="mmr", search_kwargs={"k": 5, "fetch_k": 10})
-    return retriever
+    # Applying HNSW for efficient similarity search
+    documents.embeddings.faiss.ann['HNSW'] = {'nlinks': 16}
+    
+    return documents
 
 
 @st.cache_resource