PubmedSearch / app.py
liang-huggingface's picture
Create app.py
7c7d36b
raw
history blame
5.27 kB
import gradio as gr
import pandas as pd
from Bio import Entrez
import requests
import os
HF_API = os.getenv('HF_API')
# Function to search PubMed for articles
def search_pubmed(query, retmax):
Entrez.email = 'example@example.com'
handle = Entrez.esearch(db="pubmed", term=query, retmax=retmax)
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()
article_list = []
for article in articles:
article_dict = {
'PMID': str(article['MedlineCitation']['PMID']),
'Authors': ' '.join([author['LastName'] + ' ' + author.get('Initials', '')
for author in article['MedlineCitation']['Article']['AuthorList']]),
'Title': article['MedlineCitation']['Article']['ArticleTitle'],
'Abstract': article['MedlineCitation']['Article'].get('Abstract', {}).get('AbstractText', [None])[0]
}
article_list.append(article_dict)
return pd.DataFrame(article_list)
# Function to summarize articles using Hugging Face's API
def summarize_with_huggingface(model, selected_articles):
API_URL = f"https://api-inference.huggingface.co/models/{model}"
# Your Hugging Face API key
API_KEY = HF_API
headers = {"Authorization": f"Bearer {API_KEY}"}
# Prepare the text to summarize: concatenate all abstracts
print(type(selected_articles))
print(selected_articles.to_dict(orient='records'))
text_to_summarize = " ".join(
[f"PMID: {article['PMID']}. Authors: {article['Authors']}. Title: {article['Title']}. Abstract: {article['Abstract']}."
for article in selected_articles.to_dict(orient='records')]
)
# Define the payload
payload = {
"inputs": text_to_summarize,
"parameters": {"max_length": 300} # Adjust as needed
}
# Make the POST request to the Hugging Face API
response = requests.post(API_URL, headers=headers, json=payload)
response.raise_for_status() # Raise an HTTPError if the HTTP request returned an unsuccessful status code
# The API returns a list of dictionaries. We extract the summary from the first one.
return response.json()[0]['generated_text']
import gradio as gr
from Bio import Entrez
# Always tell NCBI who you are
Entrez.email = "your.email@example.com"
def process_query(keywords, top_k):
articles = search_pubmed(keywords, top_k)
# Convert each article from a dictionary to a list of values in the correct order
articles_for_display = [[article['pmid'], article['authors'], article['title'], article['abstract']] for article in articles]
return articles_for_display
def summarize_articles(indices, articles_for_display):
# Convert indices to a list of integers
selected_indices = [int(index.strip()) for index in indices.split(',') if index.strip().isdigit()]
# Convert the DataFrame to a list of dictionaries
articles_list = articles_for_display.to_dict(orient='records')
# Select articles based on the provided indices
selected_articles = [articles_list[index] for index in selected_indices]
# Generate the summary
summary = summarize_with_huggingface(selected_articles)
return summary
# Gradio interface
with gr.Blocks() as demo:
gr.Markdown("### PubMed Article Summarizer")
model_input = gr.Textbox(label="Enter the model to use", value="h2oai/h2ogpt-4096-llama2-7b-chat")
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")
output_table = gr.Dataframe(headers=["PMID", "Authors", "Title","Abstract" ])
summarize_button = gr.Button("Summarize")
summary_output = gr.Textbox()
def update_output_table(query, retmax):
df = search_pubmed(query, retmax)
# output_table.update(value=df)
return df
search_button.click(update_output_table, inputs=[query_input, retmax_input], outputs=output_table)
summarize_button.click(fn=summarize_with_huggingface, inputs=[model_input, output_table], outputs=summary_output)
demo.launch(debug=True)
if False:
with gr.Blocks() as demo:
gr.Markdown("### PubMed Article Summarizer")
with gr.Row():
query_input = gr.Textbox(label="Query Keywords")
top_k_input = gr.Slider(minimum=1, maximum=20, value=5, step=1, label="Top K Results")
search_button = gr.Button("Search")
output_table = gr.Dataframe(headers=["Title", "Authors", "Abstract", "PMID"])
indices_input = gr.Textbox(label="Enter indices of articles to summarize (comma-separated)")
summarize_button = gr.Button("Summarize Selected Articles")
summary_output = gr.Textbox(label="Summary")
search_button.click(
fn=process_query,
inputs=[query_input, top_k_input],
outputs=output_table
)
summarize_button.click(
fn=summarize_articles,
inputs=[indices_input, output_table],
outputs=summary_output
)
demo.launch(debug=True)