Spaces:
Sleeping
Sleeping
File size: 4,593 Bytes
9c7387c 5b8a3ff 9c7387c 5b8a3ff 9c7387c bb79d76 5b8a3ff 9c7387c bb79d76 5b8a3ff bb79d76 5b8a3ff bb79d76 5b8a3ff bb79d76 5b8a3ff bb79d76 5b8a3ff bb79d76 5b8a3ff bb79d76 5b8a3ff bb79d76 5b8a3ff bb79d76 9c7387c 5b8a3ff 9c7387c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 |
from Bio import Entrez
import os
from sentence_transformers import SentenceTransformer, util
# ---------------------------- Helper Functions ----------------------------
def log_error(message: str):
"""Logs an error message to the console and a file (if possible)."""
print(f"ERROR: {message}")
try:
with open("error_log.txt", "a") as f:
f.write(f"{message}\n")
except:
print("Couldn't write to error log file.") #If logging fails, still print to console
# ---------------------------- Functions ----------------------------
def search_pubmed(query: str) -> list:
"""Searches PubMed and returns a list of article IDs using semantic search."""
ENTREZ_EMAIL = os.environ.get("ENTREZ_EMAIL", "[email protected]")
try:
Entrez.email = ENTREZ_EMAIL
print(f"Entrez Email: {Entrez.email}") # DEBUG: Check the email being used
print(f"PubMed Query: {query}") # DEBUG: Check the query being sent
"" Added code ""
# Fetch PubMed IDs
try:
handle = Entrez.esearch(db="pubmed", term=query, retmax="5")
record = Entrez.read(handle)
id_list = record["IdList"]
print(f"Original PubMed Results: {id_list}")
except Exception as e:
log_error(f"First step of PubMed search error: {e}")
return [f"First step of Error during PubMed search: {e}"]
"" Added code ""
# Fetch abstracts for all IDs:
abstracts = []
for article_id in id_list:
abstract_text = fetch_abstract(article_id)
if "Error" in abstract_text:
print(f"Skipping article {article_id} due to abstract error: {abstract_text}")
continue
abstracts.append(abstract_text)
"" Added code ""
# Semantic Search Using Sentence Transformers:
try:
model = SentenceTransformer('all-mpnet-base-v2') #Model by all-mpnet-base-v2
# Generate embeddings for abstracts and the query:
query_embedding = model.encode(query)
abstract_embeddings = model.encode(abstracts)
# Calculate cosine similarities
similarities = util.cos_sim(query_embedding, abstract_embeddings)[0]
# Sort by similarity (higher is better)
ranked_articles = sorted(zip(id_list, similarities), key=lambda x: x[1], reverse=True)
# Extract ranked IDs:
ranked_ids = [article_id for article_id, similarity in ranked_articles]
print(f"Semantic search code: {ranked_ids}") # DEBUG: Check the results
return ranked_ids
except Exception as e:
log_error(f"Semantic search code. : {e}")
return [f"Error during semantic search code: {e}"]
except Exception as e:
log_error(f"PubMed search error: {e}")
return [f"Error during PubMed search: {e}"]
def fetch_abstract(article_id: str) -> str:
"""Fetches the abstract for a given PubMed article ID."""
ENTREZ_EMAIL = os.environ.get("ENTREZ_EMAIL", "[email protected]")
try:
Entrez.email = ENTREZ_EMAIL
handle = Entrez.efetch(db="pubmed", id=article_id, rettype="abstract", retmode="text")
abstract = handle.read()
handle.close()
return abstract
except Exception as e:
log_error(f"Error fetching abstract for {article_id}: {e}")
return f"Error fetching abstract for {article_id}: {e}"
def medai_agent(query: str) -> str:
"""Orchestrates the medical literature review and presents abstract."""
article_ids = search_pubmed(query)
if isinstance(article_ids, list) and article_ids:
results = []
for article_id in article_ids:
abstract = fetch_abstract(article_id)
if "Error" not in abstract:
results.append(f"<div class='article'>\n"
f" <h3 class='article-id'>Article ID: {article_id}</h3>\n"
f" <p class='abstract'><strong>Abstract:</strong> {abstract}</p>\n"
f"</div>\n")
else:
results.append(f"<div class='article error'>\n"
f" <h3 class='article-id'>Article ID: {article_id}</h3>\n"
f" <p class='error-message'>Error processing article: {abstract}</p>\n"
f"</div>\n")
return "\n".join(results)
else:
return f"No articles found or error occurred: {article_ids}" |