Update pubmed_utils.py

pubmed_utils.py

@@ -1,22 +1,15 @@
 import requests
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from transformers import pipeline
-import nltk
-nltk.download("punkt")
-from nltk.tokenize import sent_tokenize
-from config import PUBMED_EMAIL
+from config import PUBMED_EMAIL, CHUNK_SIZE
 
-# Summarization model
-summarizer = pipeline(
-    "summarization",
-    model="facebook/bart-large-cnn",
-    tokenizer="facebook/bart-large-cnn"
-)
+# Summarization pipeline
+summarizer = pipeline("summarization", model="facebook/bart-large-cnn")
 
 
 def search_pubmed(query, max_results=5):
     """
-    Searches PubMed for articles matching the query. Returns a list of PMIDs.
+    Search PubMed for PMIDs matching the query.
     """
     url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
     params = {
@@ -29,13 +22,12 @@ def search_pubmed(query, max_results=5):
     }
     response = requests.get(url, params=params)
     response.raise_for_status()
-    data = response.json()
-    return data.get("esearchresult", {}).get("idlist", [])
+    return response.json().get("esearchresult", {}).get("idlist", [])
 
 
 def fetch_abstract(pmid):
     """
-    Fetches an abstract for a given PubMed ID (PMID).
+    Fetch abstract for a given PubMed ID.
     """
     url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"
     params = {
@@ -53,7 +45,7 @@ def fetch_abstract(pmid):
 
 def fetch_pubmed_abstracts(pmids):
     """
-    Fetch multiple PubMed abstracts in parallel.
+    Fetch multiple PubMed abstracts concurrently.
     """
     results = {}
     with ThreadPoolExecutor(max_workers=5) as executor:
@@ -67,28 +59,26 @@ def fetch_pubmed_abstracts(pmids):
     return results
 
 
-def summarize_text(text, chunk_size=512):
+def summarize_text(text, chunk_size=CHUNK_SIZE):
     """
-    Summarizes long text by splitting it into chunks and processing each chunk.
+    Summarize long text using a chunking strategy.
     """
-    sentences = sent_tokenize(text)
+    sentences = text.split(". ")
     chunks = []
     current_chunk = []
     current_length = 0
 
     for sentence in sentences:
-        if current_length + len(sentence.split()) > chunk_size:
+        tokens = len(sentence.split())
+        if current_length + tokens > chunk_size:
             chunks.append(" ".join(current_chunk))
             current_chunk = []
             current_length = 0
         current_chunk.append(sentence)
-        current_length += len(sentence.split())
+        current_length += tokens
 
     if current_chunk:
         chunks.append(" ".join(current_chunk))
 
-    summaries = [
-        summarizer(chunk, max_length=100, min_length=30, do_sample=False)[0]["summary_text"]
-        for chunk in chunks
-    ]
+    summaries = [summarizer(chunk, max_length=100, min_length=30)[0]["summary_text"] for chunk in chunks]
     return " ".join(summaries)