diff --git a/pubmed/scripts/answer_questions_ref.py b/pubmed/scripts/answer_questions_ref.py new file mode 100644 index 0000000..898eaae --- /dev/null +++ b/pubmed/scripts/answer_questions_ref.py @@ -0,0 +1,273 @@ +"""Answer questions: vector search → expand via references → pick closest refs. + +For each question: +1. Vector-search top-k documents (default 10) +2. For each hit, get its references (pmid→pmcid), fetch their embeddings from Cosmos +3. Pick the ref_k references closest to the query vector (default 2) +4. Combine: up to k + k*ref_k documents (e.g. 10 + 20 = 30) +5. Answer with LLM + +Usage: + python answer_questions_ref.py --config config.pubmed.yaml \ + --questions complex_questions.json --output answers_refs.json \ + --top 10 --ref-k 2 +""" + +import argparse, csv, json, logging, sys, time +import numpy as np +import yaml +from azure.cosmos import CosmosClient +from azure.identity import AzureCliCredential, get_bearer_token_provider +from openai import AzureOpenAI +from tqdm import tqdm + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +log = logging.getLogger(__name__) +for n in ("azure.core.pipeline.policies.http_logging_policy", "azure.identity", "azure.cosmos"): + logging.getLogger(n).setLevel(logging.WARNING) + + +def load_pmid_to_pmcid(csv_path: str) -> dict[str, str]: + mapping = {} + with open(csv_path, newline="", encoding="utf-8") as f: + for row in csv.DictReader(f, quotechar='"'): + pmcid = (row.get("PMCID") or "").strip() + pmid = (row.get("PMID") or "").strip() + if pmcid and pmid: + mapping[pmid] = pmcid + return mapping + + +def get_embedding(text: str, embed_client: AzureOpenAI, model: str, dimensions: int, max_input_chars: int = 8000) -> list[float]: + return embed_client.embeddings.create( + input=text[:max_input_chars], model=model, dimensions=dimensions + ).data[0].embedding + + +def vector_search(container, query_vector: list[float], top: int) -> list[dict]: + sql = """ + SELECT TOP @top + c.id, c.pmcid, c.title, c.journal_title, c.abstract, c.pub_year, + c.doi, c.full_text, c.embedding, c.references, + VectorDistance(c.embedding, @embedding) AS similarity + FROM c ORDER BY VectorDistance(c.embedding, @embedding) + """ + return list(container.query_items( + query=sql, + parameters=[{"name": "@embedding", "value": query_vector}, {"name": "@top", "value": top}], + enable_cross_partition_query=True, + )) + + +def fetch_docs_by_pmcids(container, pmcids: list[str]) -> list[dict]: + if not pmcids: + return [] + placeholders = ", ".join(f"@p{i}" for i in range(len(pmcids))) + sql = f""" + SELECT c.id, c.pmcid, c.title, c.journal_title, c.abstract, c.pub_year, + c.doi, c.full_text, c.embedding + FROM c WHERE c.pmcid IN ({placeholders}) + """ + params = [{"name": f"@p{i}", "value": p} for i, p in enumerate(pmcids)] + return list(container.query_items(query=sql, parameters=params, enable_cross_partition_query=True)) + + +def refs_to_pmcids(refs: list[dict], pmid_map: dict[str, str]) -> list[str]: + result = [] + for r in refs: + pmcid = (r.get("pmcid") or "").strip() + if not pmcid: + pmid = (r.get("pmid") or "").strip() + pmcid = pmid_map.get(pmid, "") + if pmcid: + result.append(pmcid) + return result + + +def closest_select(docs: list[dict], query_vec: list[float], k: int) -> list[dict]: + """Select k docs closest to query_vec by cosine similarity.""" + if len(docs) <= k: + return docs + q = np.array(query_vec, dtype=np.float64) + q_norm = np.linalg.norm(q) + if q_norm < 1e-12: + return docs[:k] + q = q / q_norm + scored = [] + for doc in docs: + emb = doc.get("embedding") + if emb and isinstance(emb, list) and len(emb) > 0: + v = np.array(emb, dtype=np.float64) + v_norm = np.linalg.norm(v) + sim = float(np.dot(q, v) / max(v_norm, 1e-12)) + scored.append((sim, doc)) + scored.sort(key=lambda x: x[0], reverse=True) + return [doc for _, doc in scored[:k]] + + +def build_context(docs: list[dict], max_chars: int = 230000) -> str: + parts, total = [], 0 + for i, doc in enumerate(docs, 1): + title = doc.get("title", "(no title)") + pmcid = doc.get("pmcid", "?") + abstract = doc.get("abstract", "") + full_text = doc.get("full_text", "") + sim = doc.get("similarity") + sim_str = f" (similarity: {sim:.4f})" if sim is not None else "" + section = f"[{i}] PMCID: {pmcid}{sim_str}\nTitle: {title}\n" + if abstract: + section += f"Abstract: {abstract}\n" + if full_text: + remaining = max_chars - total - len(section) - 200 + if remaining > 500: + section += f"Full text excerpt: {full_text[:6000][:remaining]}\n" + if total + len(section) > max_chars: + break + parts.append(section) + total += len(section) + return "\n".join(parts) + + +def answer_question(question: str, context: str, llm_client: AzureOpenAI, + model: str, system_prompt: str, temperature: float, + max_completion_tokens: int, max_retries: int) -> str: + prompt = f"""Based on the following retrieved PubMed articles, answer this question: + +Question: {question} + +Retrieved Articles: +{context} + +Provide a comprehensive answer synthesizing the information from these articles.""" + for attempt in range(max_retries): + try: + resp = llm_client.chat.completions.create( + model=model, + messages=[{"role": "system", "content": system_prompt}, + {"role": "user", "content": prompt}], + temperature=temperature, max_completion_tokens=max_completion_tokens, + ) + return resp.choices[0].message.content.strip() + except Exception as e: + log.warning(f"LLM error (attempt {attempt+1}): {e}") + if attempt < max_retries - 1: + time.sleep(min(5.0 * (2 ** attempt), 60)) + else: + return f"Error: {e}" + + +def main(): + pa = argparse.ArgumentParser(description="Answer questions: vector + reference expansion + closest selection") + pa.add_argument("--config", default="config.pubmed.yaml", help="Path to YAML config file") + pa.add_argument("--questions", default=None) + pa.add_argument("--output", default=None) + pa.add_argument("--top", "-k", type=int, default=None, help="Initial vector search docs") + pa.add_argument("--ref-k", type=int, default=None, help="Closest reference docs per hit") + args = pa.parse_args() + + with open(args.config) as f: + cfg = yaml.safe_load(f) + cosmos_cfg = cfg.get("cosmos", {}) + llm_cfg = cfg.get("llm", {}) + embed_cfg = cfg.get("embedding", {}) + answer_cfg = cfg.get("answer", {}) + + # Resolve parameters: CLI overrides > config > defaults + questions_path = args.questions or answer_cfg.get("questions", "questions.json") + output_path = args.output or answer_cfg.get("output", "answers_refs.json") + top_k = args.top if args.top is not None else int(answer_cfg.get("top", 10)) + ref_k = args.ref_k if args.ref_k is not None else int(answer_cfg.get("ref_k", 2)) + + cosmos_uri = cosmos_cfg.get("uri", "") + database_name = cosmos_cfg.get("database_name", "pubmed") + container_name = cosmos_cfg.get("sources", [{}])[0].get("container_name", "articles") + llm_endpoint = llm_cfg.get("llm_endpoint", "") + llm_model = llm_cfg.get("llm_model", "gpt-5.4") + llm_api_version = llm_cfg.get("api_version", "2024-12-01-preview") + llm_token_scope = llm_cfg.get("token_scope", "https://cognitiveservices.azure.com/.default") + llm_temperature = float(llm_cfg.get("temperature", 0.0)) + llm_max_tokens = int(llm_cfg.get("max_completion_tokens", 4096)) + llm_max_retries = int(llm_cfg.get("max_retries", 5)) + embed_endpoint = embed_cfg.get("embed_endpoint", "") + embed_api_key = embed_cfg.get("embed_api_key", "") + embed_api_version = embed_cfg.get("api_version", "2024-12-01-preview") + embed_model = embed_cfg.get("embed_model", "text-embedding-3-small") + embed_dims = int(embed_cfg.get("embed_dimensions", 1536)) + embed_max_input_chars = int(embed_cfg.get("max_input_chars", 8000)) + pmc_ids_csv = answer_cfg.get("pmc_ids_csv", "") + max_context_chars = int(answer_cfg.get("max_context_chars", 230000)) + system_prompt = answer_cfg.get("system_prompt", "You are a biomedical research assistant.").strip() + + with open(questions_path) as f: + raw = json.load(f) + questions = [item["question"] for item in raw] if raw and isinstance(raw[0], dict) else raw + log.info(f"Loaded {len(questions)} questions") + + pmid_map = load_pmid_to_pmcid(pmc_ids_csv) if pmc_ids_csv else {} + log.info(f"Loaded {len(pmid_map)} PMID→PMCID mappings") + + cred = AzureCliCredential() + container = CosmosClient(cosmos_uri, credential=cred) \ + .get_database_client(database_name).get_container_client(container_name) + embed_client = AzureOpenAI( + azure_endpoint=embed_endpoint, api_key=embed_api_key, api_version=embed_api_version) + llm_client = AzureOpenAI( + azure_endpoint=llm_endpoint, + azure_ad_token_provider=get_bearer_token_provider(cred, llm_token_scope), + api_version=llm_api_version) + log.info(f"Strategy: {top_k} vector docs, {ref_k} closest refs each → max {top_k + top_k * ref_k} docs") + + results = [] + for i, question in enumerate(tqdm(questions, desc="Processing")): + log.info(f"[{i+1}/{len(questions)}] {question[:80]}...") + query_vec = get_embedding(question, embed_client, embed_model, embed_dims, embed_max_input_chars) + + # 1. Vector search + hits = vector_search(container, query_vec, top=top_k) + seen_pmcids = {d.get("pmcid") for d in hits} + log.info(f" Retrieved {len(hits)} hits") + + # 2. For each hit, get reference PMCIDs, fetch them, pick ref_k closest to query + ref_docs_all = [] + if ref_k > 0: + for doc in hits: + refs = doc.get("references") or [] + ref_pmcids = [p for p in refs_to_pmcids(refs, pmid_map) if p not in seen_pmcids] + if not ref_pmcids: + continue + fetched = fetch_docs_by_pmcids(container, ref_pmcids) + if not fetched: + continue + selected = closest_select(fetched, query_vec, ref_k) + for s in selected: + if s.get("pmcid") not in seen_pmcids: + seen_pmcids.add(s.get("pmcid")) + ref_docs_all.append(s) + + all_docs = hits + ref_docs_all + log.info(f" Total docs: {len(hits)} hits + {len(ref_docs_all)} ref expansions = {len(all_docs)}") + + # 3. Answer + context = build_context(all_docs, max_chars=max_context_chars) + answer = answer_question(question, context, llm_client, + model=llm_model, system_prompt=system_prompt, + temperature=llm_temperature, + max_completion_tokens=llm_max_tokens, + max_retries=llm_max_retries) + + # 4. Collect + retrieved = [{"pmcid": d.get("pmcid"), "title": d.get("title"), + "journal": d.get("journal_title"), "year": d.get("pub_year"), + "similarity": d.get("similarity"), "source": "vector" if d in hits else "reference"} + for d in all_docs] + results.append({"question": question, "answer": answer, "retrieved_documents": retrieved, + "vector_count": len(hits), "ref_count": len(ref_docs_all)}) + + with open(output_path, "w") as f: + json.dump(results, f, indent=2, ensure_ascii=False) + + log.info(f"Done. Saved {len(results)} answers to {output_path}") + + +if __name__ == "__main__": + main() diff --git a/pubmed/scripts/config.pubmed.yaml b/pubmed/scripts/config.pubmed.yaml index b8f8311..2f20719 100644 --- a/pubmed/scripts/config.pubmed.yaml +++ b/pubmed/scripts/config.pubmed.yaml @@ -54,11 +54,11 @@ paths: # --------------------------------------------------------------------------- llm: llm_endpoint: "" - api_version: "2024-05-01-preview" - llm_model: "gpt-5.4-mini" + api_version: "2024-12-01-preview" + llm_model: "" temperature: 0.0 max_completion_tokens: 4096 - max_retries: 2 + max_retries: 5 premium_max_concurrency: 4 prompt_cache_size: 2048 use_rbac_auth: true @@ -79,10 +79,11 @@ local_llm: # --------------------------------------------------------------------------- embedding: embed_endpoint: "" - api_version: "2024-05-01-preview" + api_version: "2024-12-01-preview" embed_model: "text-embedding-3-small" embed_dimensions: 1536 embed_cache_size: 4096 + max_input_chars: 8000 use_rbac_auth: true embed_api_key: "" token_scope: "https://cognitiveservices.azure.com/.default" @@ -197,3 +198,14 @@ cosmos: { "path": "/full_text", "language": "en-US" } ] } + +# --------------------------------------------------------------------------- +# Answer generation settings +# --------------------------------------------------------------------------- +answer: + pmc_ids_csv: "data/PMC-ids.csv" + max_context_chars: 230000 + system_prompt: | + You are a biomedical research assistant. You are given a question and a set of retrieved PubMed articles (title, abstract, and potentially full text excerpts). + + Synthesize the information from these articles to provide a comprehensive, evidence-based answer to the question. Cite the papers by their PMCID when referencing specific findings. If the retrieved articles do not contain sufficient information to fully answer the question, state what is covered and what gaps remain. \ No newline at end of file diff --git a/pubmed/scripts/evaluate_answers_ragas.py b/pubmed/scripts/evaluate_answers_ragas.py new file mode 100644 index 0000000..c380e0e --- /dev/null +++ b/pubmed/scripts/evaluate_answers_ragas.py @@ -0,0 +1,219 @@ +"""Compare two answer JSON files using RAGAS metrics. + +Evaluates each answer set independently on multiple criteria, then prints +a side-by-side comparison. Uses Azure OpenAI via litellm + Azure AD token. + +Usage: + python evaluate_answers_ragas.py \ + --base output/review_answers_64.json \ + --other output/review_answers_divdet_64_128_rp5_eta01.json \ + --base-name vector_64 \ + --other-name divdet_64 \ + --output output/eval_ragas_vector_64_vs_divdet_64.json +""" + +import argparse +import asyncio +import json +import logging +import sys +import time + +import litellm +import yaml +from azure.identity import AzureCliCredential +from ragas.llms import llm_factory +from ragas.metrics import NumericMetric +from tqdm import tqdm + +logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") +log = logging.getLogger(__name__) + +# Suppress noisy loggers +for name in ["azure", "httpx", "openai", "litellm", "instructor", "httpcore"]: + logging.getLogger(name).setLevel(logging.WARNING) + +# --------------------------------------------------------------------------- +# Metrics definitions +# --------------------------------------------------------------------------- +METRICS = { + "comprehensiveness": NumericMetric( + name="comprehensiveness", + prompt=( + "Rate how comprehensively the answer covers the question on a scale " + "of 0.0 to 1.0. Consider breadth of coverage, depth of detail, and " + "whether important aspects are missing.\n\n" + "Question: {question}\nAnswer: {answer}" + ), + ), + "evidence_use": NumericMetric( + name="evidence_use", + prompt=( + "Rate how well the answer uses evidence from retrieved sources on a " + "scale of 0.0 to 1.0. Consider citation of specific papers, use of " + "concrete data/findings, and grounding of claims.\n\n" + "Question: {question}\nAnswer: {answer}" + ), + ), + "accuracy": NumericMetric( + name="accuracy", + prompt=( + "Rate the factual accuracy and scientific rigor of the answer on a " + "scale of 0.0 to 1.0. Consider whether claims are well-supported, " + "nuanced, and free of contradictions or errors.\n\n" + "Question: {question}\nAnswer: {answer}" + ), + ), + "clarity": NumericMetric( + name="clarity", + prompt=( + "Rate the clarity and organization of the answer on a scale of 0.0 " + "to 1.0. Consider logical structure, readability, effective use of " + "headings/sections, and overall coherence.\n\n" + "Question: {question}\nAnswer: {answer}" + ), + ), + "relevance": NumericMetric( + name="relevance", + prompt=( + "Rate how relevant the answer is to the specific question asked on " + "a scale of 0.0 to 1.0. Consider whether it directly addresses the " + "question without excessive tangential content.\n\n" + "Question: {question}\nAnswer: {answer}" + ), + ), +} + + +def _make_llm(cred, endpoint, deployment, api_version, token_scope): + """Create a fresh ragas LLM with a new token.""" + token = cred.get_token(token_scope).token + return llm_factory( + f"azure/{deployment}", + provider="litellm", + client=litellm.completion, + api_base=endpoint, + api_version=api_version, + api_key=token, + ) + + +def build_llm(llm_cfg): + """Create a ragas LLM using Azure AD token via litellm.""" + endpoint = llm_cfg.get("llm_endpoint", "") + deployment = llm_cfg.get("llm_model", "gpt-5.4") + api_version = llm_cfg.get("api_version", "2024-12-01-preview") + token_scope = llm_cfg.get("token_scope", "https://cognitiveservices.azure.com/.default") + cred = AzureCliCredential() + llm = _make_llm(cred, endpoint, deployment, api_version, token_scope) + return llm, cred, {"endpoint": endpoint, "deployment": deployment, + "api_version": api_version, "token_scope": token_scope} + + +def load_answers(path: str) -> list[dict]: + with open(path) as f: + return json.load(f) + + +def evaluate_answers(answers: list[dict], llm, cred, llm_params: dict, name: str) -> list[dict]: + """Score each answer on all metrics. Returns list of per-question result dicts.""" + results = [] + for i, item in enumerate(tqdm(answers, desc=f"Scoring [{name}]")): + question = item.get("question", "") + answer = item.get("answer", "") + row = {"question_idx": i, "question": question[:120]} + + # Rebuild LLM with fresh token every 10 questions to avoid expiry + if i % 10 == 0 and i > 0: + llm = _make_llm(cred, **llm_params) + log.info("Rebuilt LLM with fresh token") + + for metric_name, metric in METRICS.items(): + for attempt in range(5): + try: + result = metric.score(llm=llm, question=question, answer=answer) + row[metric_name] = result.value + break + except Exception as e: + err_str = str(e) + if "Authentication" in err_str or "401" in err_str or "403" in err_str: + log.warning(f"[{name}] q{i} {metric_name} auth error, rebuilding LLM...") + llm = _make_llm(cred, **llm_params) + if attempt < 4: + wait = min(5.0 * (2 ** attempt), 60) + log.warning(f"[{name}] q{i} {metric_name} attempt {attempt+1} failed: {err_str[:100]}. Retrying in {wait:.0f}s...") + time.sleep(wait) + else: + log.error(f"[{name}] q{i} {metric_name} failed after 5 attempts: {e}") + row[metric_name] = None + + results.append(row) + if (i + 1) % 10 == 0: + log.info(f"[{name}] Scored {i+1}/{len(answers)} questions") + + log.info(f"[{name}] Done scoring {len(answers)} questions") + return results + + +def main(): + parser = argparse.ArgumentParser(description="Compare two answer files using RAGAS metrics") + parser.add_argument("--config", default="config.pubmed.yaml", help="Path to YAML config file") + parser.add_argument("--base", required=True, help="Base answer JSON file") + parser.add_argument("--other", required=True, help="Other answer JSON file") + parser.add_argument("--base-name", default="base", help="Name for base method") + parser.add_argument("--other-name", default="other", help="Name for other method") + parser.add_argument("--output", "-o", default=None, help="Output JSON file") + args = parser.parse_args() + + with open(args.config) as f: + cfg = yaml.safe_load(f) + llm_cfg = cfg.get("llm", {}) + + base_answers = load_answers(args.base) + other_answers = load_answers(args.other) + log.info(f"Loaded {len(base_answers)} from {args.base}, {len(other_answers)} from {args.other}") + + assert len(base_answers) == len(other_answers), "Answer files must have the same number of questions" + + llm, cred, llm_params = build_llm(llm_cfg) + + log.info(f"Evaluating '{args.base_name}'...") + base_results = evaluate_answers(base_answers, llm, cred, llm_params, args.base_name) + + log.info(f"Evaluating '{args.other_name}'...") + other_results = evaluate_answers(other_answers, llm, cred, llm_params, args.other_name) + + # Combine results + combined = [] + for b, o in zip(base_results, other_results): + row = {"question_idx": b["question_idx"], "question": b["question"]} + for metric_name in METRICS: + row[f"{args.base_name}_{metric_name}"] = b.get(metric_name) + row[f"{args.other_name}_{metric_name}"] = o.get(metric_name) + combined.append(row) + + # Save + output_path = args.output or f"eval_ragas_{args.base_name}_vs_{args.other_name}.json" + with open(output_path, "w") as f: + json.dump(combined, f, indent=2, ensure_ascii=False) + log.info(f"Saved results to {output_path}") + + # Print summary + print(f"\n{'='*70}") + print(f"RAGAS Comparison: {args.base_name} vs {args.other_name}") + print(f"{'='*70}") + print(f"{'Metric':<25} {args.base_name:>15} {args.other_name:>15} {'delta':>10}") + print(f"{'-'*70}") + for metric_name in METRICS: + base_vals = [r.get(f"{args.base_name}_{metric_name}") for r in combined if r.get(f"{args.base_name}_{metric_name}") is not None] + other_vals = [r.get(f"{args.other_name}_{metric_name}") for r in combined if r.get(f"{args.other_name}_{metric_name}") is not None] + base_mean = sum(base_vals) / len(base_vals) if base_vals else 0 + other_mean = sum(other_vals) / len(other_vals) if other_vals else 0 + delta = other_mean - base_mean + sign = "+" if delta >= 0 else "" + print(f"{metric_name:<25} {base_mean:>15.4f} {other_mean:>15.4f} {sign}{delta:>9.4f}") + print(f"{'='*70}") + + +if __name__ == "__main__": + main() diff --git a/pubmed/scripts/readme.md b/pubmed/scripts/readme.md new file mode 100644 index 0000000..2dfdf9f --- /dev/null +++ b/pubmed/scripts/readme.md @@ -0,0 +1,37 @@ +Run script without using graph augmentation (references): + +``` +python answer_questions_ref.py --config config.pubmed.yaml --questions ../data/complex_questions.json --output output/answers/answers_top15.json --top 15 --ref-k 0 +``` + +This script will produce an answer file where each answer is generated from 15 documents. + +Run script with graph augmentation (references): + +``` +python answer_questions_ref.py --config config.pubmed.yaml --questions ../data/complex_questions.json --output output/answers_ref/answers_top5_refk4.json --top 5 --ref-k 4 +``` + +The last script will produce an answer file where each answer is generated from 16.78 documents on average (the average over 100 questions). + +Now we can compare the 2 produced answer files using ragas: + +``` +python evaluate_answers_ragas.py --base output/answers/answers_top15.json --other output/answers_ref/answers_top5_refk4.json --base-name vector15 --other-name refs_5_4 --output output/eval_ragas.json +``` + +This should produce comparison similar to this: + +
+===================================================================== +RAGAS Comparison: vector15 vs refs_5_4 +====================================================================== +Metric vector15 refs_5_4 delta +---------------------------------------------------------------------- +comprehensiveness 0.9278 0.9240 -0.0038 +evidence_use 0.8940 0.9060 + 0.0120 +accuracy 0.8908 0.8899 -0.0009 +clarity 0.9329 0.9349 + 0.0020 +relevance 0.9429 0.9414 -0.0015 +====================================================================== +\ No newline at end of file