mgbam commited on
Commit
1b899cb
·
1 Parent(s): c66f0c9

Add application file

Browse files
Files changed (1) hide show
  1. app.py +114 -0
app.py ADDED
@@ -0,0 +1,114 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from Bio import Entrez
3
+ from transformers import pipeline
4
+ import spacy
5
+
6
+ # ---------------------------- Configuration ----------------------------
7
+ ENTREZ_EMAIL = "[email protected]" # Replace with your email!
8
+ HUGGINGFACE_API_TOKEN = "HUGGINGFACE_API_TOKEN" # Replace with your Hugging Face API token!
9
+ SUMMARIZATION_MODEL = "facebook/bart-large-cnn" # Or try "google/pegasus-large"
10
+ SPACY_MODEL = "en_core_web_sm" # Small model for faster processing
11
+ # ---------------------------- Global Variables ----------------------------
12
+
13
+ summarizer = None # Initialized in the setup function
14
+ nlp = None # Initialized in the setup function
15
+
16
+ # ---------------------------- Tool Functions ----------------------------
17
+
18
+ def search_pubmed(query: str) -> list:
19
+ """Searches PubMed and returns a list of article IDs."""
20
+ Entrez.email = ENTREZ_EMAIL
21
+ try:
22
+ handle = Entrez.esearch(db="pubmed", term=query, retmax="5") # Limit to 5 for demonstration
23
+ record = Entrez.read(handle)
24
+ handle.close()
25
+ return record["IdList"]
26
+ except Exception as e:
27
+ return [f"Error during PubMed search: {e}"]
28
+
29
+ def fetch_abstract(article_id: str) -> str:
30
+ """Fetches the abstract for a given PubMed article ID."""
31
+ Entrez.email = ENTREZ_EMAIL
32
+ try:
33
+ handle = Entrez.efetch(db="pubmed", id=article_id, rettype="abstract", retmode="text")
34
+ abstract = handle.read()
35
+ handle.close()
36
+ return abstract
37
+ except Exception as e:
38
+ return f"Error fetching abstract for {article_id}: {e}"
39
+
40
+ def summarize_abstract(abstract: str) -> str:
41
+ """Summarizes an abstract using a transformer model."""
42
+ global summarizer # Access the global summarizer
43
+ if summarizer is None:
44
+ return "Summarizer not initialized. Please reload the interface."
45
+
46
+ try:
47
+ summary = summarizer(abstract, max_length=130, min_length=30, do_sample=False)[0]['summary_text']
48
+ return summary
49
+ except Exception as e:
50
+ return f"Error during summarization: {e}"
51
+
52
+ def extract_entities(text: str) -> list:
53
+ """Extracts entities (simplified) using SpaCy."""
54
+ global nlp
55
+ try:
56
+ doc = nlp(text)
57
+ entities = [(ent.text, ent.label_) for ent in doc.ents]
58
+ return entities
59
+ except Exception as e:
60
+ return [f"Error during entity extraction: {e}"]
61
+
62
+
63
+ # ---------------------------- Agent Function (Simplified) ----------------------------
64
+
65
+ def medai_agent(query: str) -> str:
66
+ """Orchestrates the medical literature review and summarization."""
67
+ article_ids = search_pubmed(query)
68
+
69
+ if isinstance(article_ids, list) and article_ids:
70
+ results = []
71
+ for article_id in article_ids:
72
+ abstract = fetch_abstract(article_id)
73
+ if "Error" not in abstract:
74
+ summary = summarize_abstract(abstract)
75
+ entities = extract_entities(abstract) #extract_entities(abstract)
76
+ results.append(f"**Article ID:** {article_id}\n\n**Summary:** {summary}\n\n**Entities:** {entities}\n\n---\n")
77
+ else:
78
+ results.append(f"Error processing article {article_id}: {abstract}\n\n---\n")
79
+ return "\n".join(results)
80
+ else:
81
+ return f"No articles found or error occurred: {article_ids}"
82
+
83
+ # ---------------------------- Gradio Interface ----------------------------
84
+
85
+ def setup():
86
+ """Initializes the summarization model and Spacy model."""
87
+ global summarizer, nlp
88
+ try:
89
+ summarizer = pipeline("summarization", model=SUMMARIZATION_MODEL, token=HUGGINGFACE_API_TOKEN)
90
+ nlp = spacy.load(SPACY_MODEL)
91
+ return "MedAI Agent initialized successfully!"
92
+ except Exception as e:
93
+ return f"Initialization error: {e}"
94
+
95
+ def launch_gradio():
96
+ """Launches the Gradio interface."""
97
+ with gr.Blocks() as iface:
98
+ gr.Markdown("# MedAI: Medical Literature Review and Summarization")
99
+ status = gr.Textbox(value="Initializing...", interactive=False)
100
+ query_input = gr.Textbox(lines=3, placeholder="Enter your medical query (e.g., 'new treatments for diabetes')...")
101
+ submit_button = gr.Button("Submit")
102
+ output_results = gr.Markdown()
103
+
104
+ submit_button.click(medai_agent, inputs=query_input, outputs=output_results)
105
+
106
+ # Initialization setup
107
+ status.value = setup()
108
+
109
+ iface.launch()
110
+
111
+ # ---------------------------- Main Execution ----------------------------
112
+
113
+ if __name__ == "__main__":
114
+ launch_gradio()