Update app.py

app.py

@@ -17,66 +17,46 @@ from langchain_core.rate_limiters import InMemoryRateLimiter
 from urllib.parse import urljoin
 
 
-def initialize(arxivcode):  
-    #loader = ArxivLoader(query=str(arxivcode),)
-    #docs = loader.load()
-    #retriever = ArxivRetriever(
-    #    load_max_docs=2,
-    #    get_full_documents=True,
-    #)
-    #docs = retriever.invoke(str(arxivcode))
-    #for i in range(len(docs)): 
-    #    docs[i].metadata['Published'] = str(docs[i].metadata['Published'])
+# LLM model
+rate_limiter = InMemoryRateLimiter(
+    requests_per_second=0.1,  # <-- MistralAI free. We can only make a request once every second
+    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.
+)    
+llm = ChatMistralAI(model="mistral-large-latest", rate_limiter=rate_limiter)
 
-    # Load, chunk and index the contents of the blog.
-    url = ['https://arxiv.org/abs/%s' % arxivcode]
-    loader = WebBaseLoader(url)
-    docs = loader.load()
+# Embeddings
+embed_model = "sentence-transformers/multi-qa-distilbert-cos-v1"
+# embed_model = "nvidia/NV-Embed-v2"
+embeddings = HuggingFaceInstructEmbeddings(model_name=embed_model)
+# embeddings = MistralAIEmbeddings()
 
+def RAG(llm, docs, embeddings):
 
-    # LLM model
-    rate_limiter = InMemoryRateLimiter(
-        requests_per_second=0.1,  # <-- MistralAI free. We can only make a request once every second
-        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.
-    )    
-    llm = ChatMistralAI(model="mistral-large-latest", rate_limiter=rate_limiter)
-    
-    # Embeddings
-    embed_model = "sentence-transformers/multi-qa-distilbert-cos-v1"
-    # embed_model = "nvidia/NV-Embed-v2"
-    embeddings = HuggingFaceInstructEmbeddings(model_name=embed_model)
-    # embeddings = MistralAIEmbeddings()
+    # Split text
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    splits = text_splitter.split_documents(docs)
 
-    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
+    # 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")
 
-    return RAG(llm, docs, embeddings)
+    # Create the chain
+    rag_chain = (
+        {"context": retriever | format_docs, "question": RunnablePassthrough()}
+        | prompt
+        | llm
+        | StrOutputParser()
+    )
+    return rag_chain
+
+def format_docs(docs):
+    return "\n\n".join(doc.page_content for doc in docs)
 
 def handle_prompt(message, history, arxivcode, rag_chain): 
     try:
@@ -91,14 +71,30 @@ def handle_prompt(message, history, arxivcode, rag_chain):
 
 greetingsmessage = "Hi, I'm your personal arXiv reader. Ask me questions about the arXiv paper above"
 
+
 with gr.Blocks() as demo:     
-  arxiv_code = gr.Textbox("", label="arxiv.number")
-  rag_chain = initialize(arxiv_code)
+    arxiv_code = gr.Textbox("", label="arxiv.number")
+    
+    #rag_chain = initialize(arxiv_code)
+    loader = ArxivLoader(query=str(arxivcode),)
+    docs = loader.load()
+    #retriever = ArxivRetriever(
+    #    load_max_docs=2,
+    #    get_full_documents=True,
+    #)
+    #docs = retriever.invoke(str(arxivcode))
+    #for i in range(len(docs)): 
+    #    docs[i].metadata['Published'] = str(docs[i].metadata['Published'])
+    
+    # Load, chunk and index the contents of the blog.
+    #url = ['https://arxiv.org/abs/%s' % arxivcode]
+    #loader = WebBaseLoader(url)
+    #docs = loader.load()
+    rag_chain = RAG(llm, docs, embeddings)
     
-  gr.ChatInterface(handle_prompt, type="messages", theme=gr.themes.Soft(), 
+    gr.ChatInterface(handle_prompt, type="messages", theme=gr.themes.Soft(), 
                           description=greetingsmessage, 
                    additional_inputs=[arxiv_code, rag_chain]
                   )
                           
-if __name__=='__main__': 
-    demo.launch()
+demo.launch()