Update app.py

app.py

@@ -46,15 +46,15 @@ def process_document(doc: Document):
     metadata["Body"] = doc.page_content
     return Document(page_content=doc.metadata["Summary"], metadata=metadata)
 
-def get_data(lookback_days: float, user_query: str):
+def get_data(category: str, lookback_days: float, user_query: str):
     print("User query:", user_query)
     max_date = date.today()
     min_date = (max_date - timedelta(days=lookback_days))
-    query = f"cat:hep-th AND submittedDate:[{min_date.strftime('%Y%m%d')} TO {max_date.strftime('%Y%m%d')}]"
+    query = f"cat:{category} AND submittedDate:[{min_date.strftime('%Y%m%d')} TO {max_date.strftime('%Y%m%d')}]"
     loader = ArxivLoader(query=query, load_max_docs=LOAD_MAX_DOCS)
     docs = [process_document(doc) for doc in loader.load()]
     if len(docs) == 0:
-        return "Found no documents. Consider increasing the value for 'Articles from this many days in the past will be searched through.'."
+        return "Found no documents. Check if the category is correct or consider increasing the value for 'Articles from this many days in the past will be searched through.'."
     db = Chroma.from_documents(docs, embeddings)
     retriever = db.as_retriever()    
     relevant_docs = retriever.get_relevant_documents(user_query)
@@ -80,7 +80,7 @@ with gr.Blocks() as demo:
         with gr.Column():
             with gr.Accordion("Parameters", open=False):
                 lookback_days = gr.Number(2, label="Articles from this many days in the past will be searched through.", minimum=1, maximum=7)
-        
+            category = gr.Textbox(value="hep-th", label="Which category to search through. See https://arxiv.org/category_taxonomy for possible values.")
             input_text = gr.Textbox(placeholder="Describe your field of research in a few words")
             gr.Examples(
                 [["Supersymmetric Conformal Field Theory"], ["Black hole information paradox"]],
@@ -89,8 +89,8 @@ with gr.Blocks() as demo:
             btn = gr.Button(value="Submit")
 
         with gr.Column():
-            output = gr.Markdown()
+            output = gr.Markdown("Press 'submit' to see your results.")
     
-    btn.click(fn=get_data, inputs=[lookback_days,input_text], outputs=output)
+    btn.click(fn=get_data, inputs=[category, lookback_days,input_text], outputs=output)
 
     demo.queue().launch()