Update app.py

app.py

@@ -4,7 +4,7 @@ from langchain.document_loaders import ArxivLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import Chroma
 from langchain.embeddings import HuggingFaceEmbeddings
-# from langchain.document_loaders import Document
+from langchain.document_loaders import Document
 
 CHUNK_SIZE = 1000
 LOAD_MAX_DOCS = 5
@@ -17,22 +17,27 @@ loader = ArxivLoader(query=query, load_max_docs=LOAD_MAX_DOCS)
 
 embeddings = HuggingFaceEmbeddings()    
 
+def process_document(doc: Document):
+    metadata = doc.metadata
+    metadata["Body"] = doc.page_content
+    return Document(page_content=doc.metadata["Summary"], metadata=metadata)
+
 def get_data(user_query: str):
-    documents = loader.load()
-    # texts = text_splitter.split_documents(documents)    
-    texts = documents
-    for doc in texts:
-        doc.page_content = doc.metadata["Summary"]
-    db = Chroma.from_documents(texts, embeddings)
+    docs = loader.load()
+    docs = [process_document(doc) for doc in docs]
+    db = Chroma.from_documents(docs, embeddings)
     retriever = db.as_retriever()    
-    docs = retriever.get_relevant_documents(user_query)
-    print(docs[0].metadata)
-    return "\n\n".join([d.page_content for d in docs])
+    relevant_docs = retriever.get_relevant_documents(user_query)
+    print(relevant_docs[0].metadata)
+    output = ""
+    for doc in relevant_docs:
+        output += f"**Title: {doc["title"]}**\nAbstract: {doc["Summary"]}\n\n"
+    return output
 
 demo = gr.Interface(
     fn=get_data,
     inputs="text",
-    outputs="text",
+    outputs=gr.Markdown(),
     title="Document Filter",
     description="Enter a query to filter the list of documents."
 )