Update pubmed_utils.py

pubmed_utils.py

@@ -1,102 +1,94 @@
 import requests
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from transformers import pipeline
 import nltk
-nltk.download('punkt')
+nltk.download("punkt")
 from nltk.tokenize import sent_tokenize
+from config import PUBMED_EMAIL
 
-from transformers import pipeline
-from config import MY_PUBMED_EMAIL
-
-# Summarization pipeline for PubMed abstracts
+# Summarization model
 summarizer = pipeline(
     "summarization",
     model="facebook/bart-large-cnn",
     tokenizer="facebook/bart-large-cnn"
 )
 
-def search_pubmed(query, max_results=3):
+
+def search_pubmed(query, max_results=5):
     """
-    Searches PubMed via ESearch. Returns list of PMIDs.
+    Searches PubMed for articles matching the query. Returns a list of PMIDs.
     """
-    base_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
+    url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
     params = {
         "db": "pubmed",
         "term": query,
         "retmax": max_results,
         "retmode": "json",
-        "tool": "ElysiumRAG",
-        "email": MY_PUBMED_EMAIL
+        "tool": "MedicalAI",
+        "email": PUBMED_EMAIL,
     }
-    resp = requests.get(base_url, params=params)
-    resp.raise_for_status()
-    data = resp.json()
+    response = requests.get(url, params=params)
+    response.raise_for_status()
+    data = response.json()
     return data.get("esearchresult", {}).get("idlist", [])
 
-def fetch_one_abstract(pmid):
+
+def fetch_abstract(pmid):
     """
-    Fetches a single abstract for the given PMID.
+    Fetches an abstract for a given PubMed ID (PMID).
     """
-    base_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"
+    url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"
     params = {
         "db": "pubmed",
+        "id": pmid,
         "retmode": "text",
         "rettype": "abstract",
-        "id": pmid,
-        "tool": "ElysiumRAG",
-        "email": MY_PUBMED_EMAIL
+        "tool": "MedicalAI",
+        "email": PUBMED_EMAIL,
     }
-    resp = requests.get(base_url, params=params)
-    resp.raise_for_status()
-    raw_text = resp.text.strip() or "No abstract text found."
-    return (pmid, raw_text)
+    response = requests.get(url, params=params)
+    response.raise_for_status()
+    return response.text.strip()
+
 
 def fetch_pubmed_abstracts(pmids):
     """
-    Parallel fetching of multiple abstracts.
+    Fetch multiple PubMed abstracts in parallel.
     """
-    if not pmids:
-        return {}
-    results_map = {}
-    with ThreadPoolExecutor(max_workers=min(len(pmids), 5)) as executor:
-        future_to_pmid = {executor.submit(fetch_one_abstract, pmid): pmid for pmid in pmids}
+    results = {}
+    with ThreadPoolExecutor(max_workers=5) as executor:
+        future_to_pmid = {executor.submit(fetch_abstract, pmid): pmid for pmid in pmids}
         for future in as_completed(future_to_pmid):
             pmid = future_to_pmid[future]
             try:
-                pmid_result, text = future.result()
-                results_map[pmid_result] = text
+                results[pmid] = future.result()
             except Exception as e:
-                results_map[pmid] = f"Error: {str(e)}"
-    return results_map
+                results[pmid] = f"Error fetching PMID {pmid}: {str(e)}"
+    return results
+
 
-def chunk_and_summarize(abstract_text, chunk_size=512):
+def summarize_text(text, chunk_size=512):
     """
-    Splits large abstracts by sentences, summarizes each chunk, then concatenates.
+    Summarizes long text by splitting it into chunks and processing each chunk.
     """
-    sentences = sent_tokenize(abstract_text)
+    sentences = sent_tokenize(text)
     chunks = []
-
     current_chunk = []
     current_length = 0
-    for sent in sentences:
-        tokens_in_sent = len(sent.split())
-        if current_length + tokens_in_sent > chunk_size:
+
+    for sentence in sentences:
+        if current_length + len(sentence.split()) > chunk_size:
             chunks.append(" ".join(current_chunk))
             current_chunk = []
             current_length = 0
-        current_chunk.append(sent)
-        current_length += tokens_in_sent
+        current_chunk.append(sentence)
+        current_length += len(sentence.split())
 
     if current_chunk:
         chunks.append(" ".join(current_chunk))
 
-    summarized_pieces = []
-    for c in chunks:
-        summary_out = summarizer(
-            c,
-            max_length=100,
-            min_length=30,
-            do_sample=False
-        )
-        summarized_pieces.append(summary_out[0]['summary_text'])
-
-    return " ".join(summarized_pieces).strip()
+    summaries = [
+        summarizer(chunk, max_length=100, min_length=30, do_sample=False)[0]["summary_text"]
+        for chunk in chunks
+    ]
+    return " ".join(summaries)