Update app.py

app.py

@@ -44,9 +44,49 @@ def generate_response(prompt):
     # Decode the response
     response = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return response
+
+def search_pubmed(query, retmax=5, mindate=None, maxdate=None, datetype="pdat"):
+    Entrez.email = '[email protected]'  # Always set the Entrez.email to tell NCBI who you are
+    search_kwargs = {
+        "db": "pubmed",
+        "term": query,
+        "retmax": retmax,
+        "sort": 'relevance',
+        "datetype": datetype
+    }
+    
+    # If dates are provided, add them to the search arguments
+    if mindate:
+        search_kwargs["mindate"] = mindate
+    if maxdate:
+        search_kwargs["maxdate"] = maxdate
+
+    handle = Entrez.esearch(**search_kwargs)
+    record = Entrez.read(handle)
+    handle.close()
+    idlist = record['IdList']
+    
+    handle = Entrez.efetch(db="pubmed", id=idlist, retmode="xml")
+    articles = Entrez.read(handle)['PubmedArticle']
+    handle.close()
+    
+    # ... (the rest of your existing code to extract article information)
+    abstracts = []
+    for article in articles:
+        article_id = article['MedlineCitation']['PMID']
+        article_title = article['MedlineCitation']['Article']['ArticleTitle']
+        print(len(article['MedlineCitation']['Article'].get('Abstract', {}).get('AbstractText', [None])))
+        abstract_text = article['MedlineCitation']['Article'].get('Abstract', {}).get('AbstractText', [None])
+        
+        if isinstance(abstract_text, list):
+            # Join the list elements if abstract is a list
+            abstract_text = " ".join(abstract_text)
+        abstracts.append((article_id, article_title, abstract_text))
+
+    return pd.DataFrame(abstracts)
     
 # Function to search PubMed for articles
-def search_pubmed(query, retmax):
+def search_pubmed_old(query, retmax):
     Entrez.email = '[email protected]'
     handle = Entrez.esearch(db="pubmed", term=query, retmax=retmax, sort = 'relevance')
     record = Entrez.read(handle)
@@ -192,6 +232,10 @@ with gr.Blocks() as demo:
     with gr.Row():
         password_input = gr.Textbox(label="Enter the password")
         model_input = gr.Textbox(label="Enter the model to use", value="h2oai/h2ogpt-4096-llama2-7b-chat")
+    with gr.Row():
+        startdate = gr.Textbox(label="Starting year")
+        enddate = gr.Textbox(label="End year")
+        
     query_input = gr.Textbox(label="Query Keywords")
     retmax_input = gr.Slider(minimum=1, maximum=20, value=5, step=1, label="Number of articles")
     search_button = gr.Button("Search")
@@ -200,11 +244,11 @@ with gr.Blocks() as demo:
     summary_output = gr.Textbox()
 
 
-    def update_output_table(query, retmax):
-        df = search_pubmed(query, retmax)
+    def update_output_table(query, retmax, startdate, enddate):
+        df = search_pubmed(query, retmax, startdate, enddate)
 #        output_table.update(value=df)
         return df
-    search_button.click(update_output_table, inputs=[query_input, retmax_input], outputs=output_table)
+    search_button.click(update_output_table, inputs=[query_input, retmax_input, startdate, enddate], outputs=output_table)
     summarize_button.click(fn=summarize_with_huggingface, inputs=[model_input, output_table, password_input], outputs=summary_output)
 
 demo.launch(debug=True)