Update app.py

app.py

@@ -4,19 +4,54 @@ from langchain.document_loaders import ArxivLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import Chroma
 from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.llms import HuggingFaceHub
+# from langchain.llms import FakeListLLM
+from langchain.chains import LLMChain, StuffDocumentsChain
+from langchain.prompts import PromptTemplate
 from langchain.schema import Document
 
-CHUNK_SIZE = 1000
-LOAD_MAX_DOCS = 100
-text_splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE)
 
+LOAD_MAX_DOCS = 100
 min_date = (date.today() - timedelta(days=2)).strftime('%Y%m%d')
 max_date = date.today().strftime('%Y%m%d')
 query = f"cat:hep-th AND submittedDate:[{min_date} TO {max_date}]"
 loader = ArxivLoader(query=query, load_max_docs=LOAD_MAX_DOCS)
 
+# CHUNK_SIZE = 1000
+# text_splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE)
+
 embeddings = HuggingFaceEmbeddings()    
 
+
+document_prompt = PromptTemplate(
+  template="Title: {Title}\nContent: {page_content}",
+  input_variables=["page_content", "Title"],
+)
+prompt = PromptTemplate(
+  template=
+  """Write a personalised newsletter for a researcher. The researcher describes his work as follows:"{context}". Base the newsletter on the following articles:\n\n"{text}"\n\nNEWSLETTER:""",
+  input_variables=["context", "text"])
+
+# llm = FakeListLLM(responses=list(map(str, range(100))))
+REPO_ID = "HuggingFaceH4/starchat-beta"
+llm = HuggingFaceHub(
+    repo_id=REPO_ID,
+    model_kwargs={
+        "max_new_tokens": 1024,
+        "do_sample": True,
+        "temperature": 0.8,
+        "top_p": 0.9
+    }
+)
+
+llm_chain = LLMChain(llm=llm, prompt=prompt, verbose=True)
+stuff_chain = StuffDocumentsChain(
+  llm_chain=llm_chain,
+  document_variable_name="text",
+  document_prompt=document_prompt,
+  verbose=True,
+)
+
 def process_document(doc: Document):
     metadata = doc.metadata
     metadata["Body"] = doc.page_content
@@ -29,10 +64,11 @@ def get_data(user_query: str):
     retriever = db.as_retriever()    
     relevant_docs = retriever.get_relevant_documents(user_query)
     print(relevant_docs[0].metadata)
-    output = ""
+    articles = ""
     for doc in relevant_docs:
-        output += f"**Title: {doc.metadata['Title']}**\nAbstract: {doc.metadata['Summary']}\n\n"
-    return output
+        articles += f"**Title: {doc.metadata['Title']}**\n\nAbstract: {doc.metadata['Summary']}\n\n"
+    output = stuff_chain({"input_documents": relevant_docs, "context": user_query})
+    return f"{output["output_text"]}\n\n\n\nUsed articles:\n\n{output}"
 
 demo = gr.Interface(
     fn=get_data,