mgbam commited on
Commit
f33af5a
·
verified ·
1 Parent(s): de576e2

Update pubmed_rag.py

Browse files
Files changed (1) hide show
  1. pubmed_rag.py +9 -8
pubmed_rag.py CHANGED
@@ -23,7 +23,7 @@ def search_pubmed(query, max_results=MAX_PUBMED_RESULTS):
23
  "tool": "AdvancedMedicalAI",
24
  "email": MY_PUBMED_EMAIL,
25
  }
26
- response = requests.get(url, params=params)
27
  response.raise_for_status()
28
  data = response.json()
29
  return data.get("esearchresult", {}).get("idlist", [])
@@ -41,7 +41,7 @@ def fetch_abstract(pmid):
41
  "tool": "AdvancedMedicalAI",
42
  "email": MY_PUBMED_EMAIL,
43
  }
44
- response = requests.get(url, params=params)
45
  response.raise_for_status()
46
  return response.text.strip()
47
 
@@ -52,7 +52,8 @@ def fetch_pubmed_abstracts(pmids):
52
  results = {}
53
  for pmid in pmids:
54
  try:
55
- results[pmid] = fetch_abstract(pmid)
 
56
  except Exception as e:
57
  results[pmid] = f"Error fetching PMID {pmid}: {e}"
58
  return results
@@ -78,8 +79,8 @@ def summarize_text(text, chunk_size=SUMMARIZATION_CHUNK_SIZE):
78
  if current_chunk:
79
  chunks.append(" ".join(current_chunk))
80
 
81
- summaries = [
82
- summarizer(chunk, max_length=100, min_length=30, do_sample=False)[0]["summary_text"]
83
- for chunk in chunks
84
- ]
85
- return " ".join(summaries)
 
23
  "tool": "AdvancedMedicalAI",
24
  "email": MY_PUBMED_EMAIL,
25
  }
26
+ response = requests.get(url, params=params, timeout=10)
27
  response.raise_for_status()
28
  data = response.json()
29
  return data.get("esearchresult", {}).get("idlist", [])
 
41
  "tool": "AdvancedMedicalAI",
42
  "email": MY_PUBMED_EMAIL,
43
  }
44
+ response = requests.get(url, params=params, timeout=10)
45
  response.raise_for_status()
46
  return response.text.strip()
47
 
 
52
  results = {}
53
  for pmid in pmids:
54
  try:
55
+ abstract = fetch_abstract(pmid)
56
+ results[pmid] = abstract
57
  except Exception as e:
58
  results[pmid] = f"Error fetching PMID {pmid}: {e}"
59
  return results
 
79
  if current_chunk:
80
  chunks.append(" ".join(current_chunk))
81
 
82
+ summaries = []
83
+ for chunk in chunks:
84
+ summary = summarizer(chunk, max_length=100, min_length=30, do_sample=False)[0]["summary_text"]
85
+ summaries.append(summary)
86
+ return " ".join(summaries)