Update app.py

app.py

@@ -22,62 +22,12 @@ rate_limiter = InMemoryRateLimiter(
     check_every_n_seconds=0.01,  # Wake up every 100 ms to check whether allowed to make a request,
     max_bucket_size=10,  # Controls the maximum burst size.
 )
-"""
-# get data
-urlsfile = open("urls.txt")
-urls = urlsfile.readlines()
-urls = [url.replace("\n","") for url in urls]
-urlsfile.close()
-
-# Load, chunk and index the contents of the blog.
-loader = WebBaseLoader(urls)
-docs = loader.load()
-
-# load arxiv papers 
-arxivfile = open("arxiv.txt")
-arxivs = arxivfile.readlines()
-arxivs = [arxiv.replace("\n","") for arxiv in arxivs]
-arxivfile.close()
 
 retriever = ArxivRetriever(
     load_max_docs=2,
     get_ful_documents=True,
 )
 
-for arxiv in arxivs: 
-    doc = retriever.invoke(arxiv) 
-    doc[0].metadata['Published'] = str(doc[0].metadata['Published'])
-    docs.append(doc[0])
-    
-
-def format_docs(docs):
-    return "\n\n".join(doc.page_content for doc in docs)
-
-def RAG(llm, docs, embeddings):
-
-    # Split text
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    splits = text_splitter.split_documents(docs)
-
-    # Create vector store
-    vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)
-
-    # Retrieve and generate using the relevant snippets of the documents
-    retriever = vectorstore.as_retriever()
-
-    # Prompt basis example for RAG systems
-    prompt = hub.pull("rlm/rag-prompt")
-
-    # Create the chain
-    rag_chain = (
-        {"context": retriever | format_docs, "question": RunnablePassthrough()}
-        | prompt
-        | llm
-        | StrOutputParser()
-    )
-
-    return rag_chain
-
 # LLM model
 llm = ChatMistralAI(model="mistral-large-latest", rate_limiter=rate_limiter)
 
@@ -87,10 +37,48 @@ embed_model = "sentence-transformers/multi-qa-distilbert-cos-v1"
 embeddings = HuggingFaceInstructEmbeddings(model_name=embed_model)
 # embeddings = MistralAIEmbeddings()
 
-# RAG chain
-rag_chain = RAG(llm, docs, embeddings)
+def initialize(arxivcode):  
+    docs = retriever.invoke(arxiv) 
+    docs[0].metadata['Published'] = str(doc[0].metadata['Published'])
+
+    def format_docs(docs):
+        return "\n\n".join(doc.page_content for doc in docs)
+    
+    def RAG(llm, docs, embeddings):
+    
+        # Split text
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+        splits = text_splitter.split_documents(docs)
+    
+        # Create vector store
+        vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)
+    
+        # Retrieve and generate using the relevant snippets of the documents
+        retriever = vectorstore.as_retriever()
+    
+        # Prompt basis example for RAG systems
+        prompt = hub.pull("rlm/rag-prompt")
+    
+        # Create the chain
+        rag_chain = (
+            {"context": retriever | format_docs, "question": RunnablePassthrough()}
+            | prompt
+            | llm
+            | StrOutputParser()
+        )
+    
+        return rag_chain
+
+    return RAG(llm, docs, embeddings)
+
+rag_chain = None 
 
-def handle_prompt(message, history):
+def handle_prompt(message, history, arxivcode): 
+    if rag_chain is None: 
+        # initialize RAG chain 
+        # RAG chain
+        rag_chain = initialize(arxivcode)
+        
     try:
         # Stream output
         out=""
@@ -99,17 +87,19 @@ def handle_prompt(message, history):
             yield out
     except:
         raise gr.Error("Requests rate limit exceeded")
-"""
 
-def handle_prompt(message, history, input1): 
-    return f"arxiv code: {input1}, {message}" 
 
+greetingsmessage = "Hi, I'm your personal arXiv reader. Ask me questions about the arXiv paper above"
 
-greetingsmessage = "Hi, I'm your personal arXiv reader. Input the arXiv number of the paper:"
+with gr.Blocks() as demo:     
 
-demo = gr.ChatInterface(handle_prompt, type="messages", theme=gr.themes.Soft(), 
-                        description=greetingsmessage, 
-                        additional_inputs=[gr.Textbox("", label="arxiv.code")]
-                       )
+  arxiv_code = gr.Textbox("", label="arxiv.number")
+  
+  gr.ChatInterface(handle_prompt, type="messages", theme=gr.themes.Soft(), 
+                          description=greetingsmessage, 
+                   additional_inputs=[arxiv_code]
+                  )
+                          
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()