PerturbReason/PerturbReason_dataset_code
012
1"""2SUMMER Pipeline with vLLM (v2 - 3way Dataset)3Uses the same 3way dataset as the Qwen baseline (dataset_0410_disambig).4Enhances prompts with knowledge-graph-derived gene descriptions.5 6Aligned with the official PerturbQA SUMMER implementation:7 https://github.com/Genentech/PerturbQA/tree/main/examples/summer8 9Input: dataset_0410_disambig/{context_type}/{split}/*.jsonl (task=3way)10Output: {source_file, prompt, ground_truth_response, model_output} (matches Qwen)11"""12import argparse13import json14import os15import re16import glob17import sys18from tqdm import tqdm19from vllm import LLM, SamplingParams20 21sys.path.append(os.path.dirname(os.path.abspath(__file__)))22try:23 import summer_pipeline24except ImportError:25 print("ERROR: summer_pipeline.py not found in the same directory.")26 sys.exit(1)27 28# ==============================================================================29# PROMPT TEMPLATES30# ==============================================================================31 32PROMPT_SUMMARIZE_GENE = """You are an expert molecular biologist who studies how genes are related using Perturb-seq.33 34Task: You are writing a brief overview of the human gene {name}, with a focus on its molecular and cellular functions. You will be provided a set of database entries about the gene. Ensure that your overview remains faithful to this domain knowledge.35 36Format:37- Write one to two sentences describing the primary molecular and cellular function of gene {name}.38- Write one sentence describing what types of perturbations might impact the expression of gene {name}. For example, you might consider pathways that are upstream of the gene or compensatory mechanisms.39 40Constraints:41- Maintain a professional tone throughout.42- Do not comment on your own writing.43- Do not add any notes or references. Do not make up additional information.44- Do not discuss the importance or impact of the gene. Focus only on its function.45 46Domain knowledge:47{entries}48 49Brief overview of gene {name}:50"""51 52# --- Genetic 3-way (up / down / unchanged) ---53PROMPT_GENETIC_3WAY = """You are an expert molecular biologist who studies how genes are related using Perturb-seq. Your goal is to determine: Is the expression of {gene} up, down, or unchanged after CRISPRi knockdown of {pert} in {cell_short}?54 55You are given as input:56- Description of perturbed gene ({pert}): description of gene that is perturbed via CRISPRi knockdown57- Description of gene of interest ({gene}): description of gene, the impact on which you wish to infer58- Context: description of cell line in which the genes are expressed59- Basal expression levels: baseline expression levels of genes in the pathway60- Reasoning paths: known molecular relationships between the perturbed gene and the gene of interest61 62Output: Please analyze the following aspects step by step.63 641) What is the function of {pert}? What downstream pathways or genes might be affected when {pert} is knocked down?652) What is the function of {gene}? What upstream pathways regulate its expression?663) Based on the reasoning paths and basal expression levels, trace the relationship from {pert} to {gene}. Consider:67 - If {pert} activates a pathway that activates {gene}, knockdown would likely decrease {gene} (down)68 - If {pert} activates a pathway that represses {gene}, knockdown would likely increase {gene} (up)69 - If there is no clear or strong causal relationship, expression would likely be unchanged704) Is the expression of {gene} up, down, or unchanged after CRISPRi knockdown of {pert} in {cell_short}?71 72Structure your response in exactly three blocks:731. <thinking>Your step-by-step biological reasoning here.</thinking>742. <answer>up</answer> OR <answer>down</answer> OR <answer>unchanged</answer>753. <triplet>A JSON list of causal edges from your reasoning, e.g., [["A", "inhibits", "B"], ["B", "activates", "C"]].</triplet>76 77[Start of Input]78- Description of perturbed gene ({pert}): {pert_desc}79- Description of gene of interest ({gene}): {gene_desc}80- Context: {cell_desc}81- Basal expression levels:82{basal_text}83- Reasoning paths:84{reasoning_paths}85[End of Input]"""86 87# --- Chemical 3-way (up / down / unchanged) ---88PROMPT_CHEMICAL_3WAY = """You are an expert molecular biologist who studies how drugs and chemical perturbations affect gene expression. Your goal is to determine: Is the expression of {gene} up, down, or unchanged after chemical perturbation {pert} in {cell_short}?89 90You are given as input:91- Description of drug ({pert}): mechanism of action and known targets92- Description of gene of interest ({gene}): molecular and cellular functions93- Context: description of cell line in which the experiment is conducted94- Basal expression levels: baseline expression levels of genes in the pathway95- Reasoning paths: known molecular relationships between the drug and the gene96 97Output: Please analyze the following aspects step by step.98 991) What are the primary targets of {pert}? What is its mechanism of action?1002) What is the function of {gene}? What upstream pathways regulate its expression?1013) Based on the reasoning paths and basal expression levels, trace the relationship from {pert} to {gene}. Consider:102 - If {pert} inhibits a target that activates {gene}, expression would likely decrease (down)103 - If {pert} inhibits a target that represses {gene}, expression would likely increase (up)104 - If there is no clear or strong causal relationship, expression would likely be unchanged1054) Is the expression of {gene} up, down, or unchanged after chemical perturbation {pert} in {cell_short}?106 107Structure your response in exactly three blocks:1081. <thinking>Your step-by-step biological reasoning here.</thinking>1092. <answer>up</answer> OR <answer>down</answer> OR <answer>unchanged</answer>1103. <triplet>A JSON list of causal edges from your reasoning, e.g., [["A", "inhibits", "B"], ["B", "activates", "C"]].</triplet>111 112[Start of Input]113- Description of drug ({pert}): {pert_desc}114- Description of gene of interest ({gene}): {gene_desc}115- Context: {cell_desc}116- Basal expression levels:117{basal_text}118- Reasoning paths:119{reasoning_paths}120[End of Input]"""121 122# Cell line descriptions (from official SUMMER)123CELL_LINES = {124 "K562 cells": "K562 cells are immortalised myelogenous leukemia cells of the erythroleukemia type.",125 "K562": "K562 cells are immortalised myelogenous leukemia cells of the erythroleukemia type.",126 "k562": "K562 cells are immortalised myelogenous leukemia cells of the erythroleukemia type.",127 "RPE1 cells": "RPE1 cells are a non-cancerous, hTERT-immortalized, near-euploid, adherent, and p53-positive cell line.",128 "rpe1": "RPE1 cells are a non-cancerous, hTERT-immortalized, near-euploid, adherent, and p53-positive cell line.",129 "Jurkat cells": "Jurkat cells are an immortalized line of human T lymphocyte cells.",130 "jurkat": "Jurkat cells are an immortalized line of human T lymphocyte cells.",131 "HepG2 cells": "HepG2 cells are a human liver cancer cell line, derived from a patient with a well-differentiated hepatocellular carcinoma.",132 "hepg2": "HepG2 cells are a human liver cancer cell line, derived from a patient with a well-differentiated hepatocellular carcinoma.",133 "HepG2": "HepG2 cells are a human liver cancer cell line, derived from a patient with a well-differentiated hepatocellular carcinoma.",134 "B cells": "B cells are lymphocytes that develop in the bone marrow and play a central role in the adaptive immune response by producing antibodies.",135}136 137# ==============================================================================138# Helper Functions139# ==============================================================================140 141def get_gene_summary_prompt(gene: str, resources: tuple) -> str:142 """Build a summarization prompt for a gene using KG data."""143 ensembl_dict, uniprot_dict, go_desc, all_graphs = resources144 kg_desc = summer_pipeline.get_gene_descriptions(145 gene, ensembl_dict, uniprot_dict, go_desc, all_graphs146 )147 if len(kg_desc) == 0:148 return None149 MAX_CHARS = 24000150 current_chars = 0151 truncated_desc = []152 for entry in kg_desc:153 entry_len = len(entry)154 if current_chars + entry_len > MAX_CHARS:155 break156 truncated_desc.append(f"- {entry}")157 current_chars += entry_len158 entries = "\n".join(truncated_desc)159 return PROMPT_SUMMARIZE_GENE.format(name=gene, entries=entries)160 161 162def format_reasoning_paths(triplets):163 """Convert triplets to reasoning path text."""164 if not triplets:165 return " No reasoning paths available."166 paths = []167 seen = set()168 for t in triplets:169 s = f"{t[0]} {t[1]} {t[2]}"170 if s not in seen:171 paths.append(f" - {s}")172 seen.add(s)173 return "\n".join(paths)174 175 176def format_basal_context(basal_context):177 """Convert basal expression dict {gene: float} to text like the dataset prompt."""178 if not basal_context:179 return " No basal expression data available."180 lines = []181 for gene, val in sorted(basal_context.items(), key=lambda x: -x[1]):182 if val >= 0.6:183 level = "high"184 elif val >= 0.2:185 level = "medium"186 else:187 level = "low"188 lines.append(f" - {gene}: {level}")189 return "\n".join(lines)190 191 192def get_pert_description_chemical(item):193 """Generate a description for a chemical perturbation from triplets."""194 pert = item["input"]["pert"]195 triplets = item["output"].get("triplets_cell_conditioned",196 item["output"].get("triplets_all", []))197 targets = set()198 mechanisms = []199 for triplet in triplets:200 if triplet[0] == pert:201 targets.add(triplet[2])202 mechanisms.append(f"{pert} {triplet[1]} {triplet[2]}")203 if targets:204 desc = f"{pert} targets {', '.join(list(targets)[:5])}."205 if mechanisms:206 desc += f" Known interactions: {'; '.join(mechanisms[:3])}."207 else:208 desc = f"{pert} is a chemical compound used for perturbation studies."209 smiles = item["input"].get("smiles", "")210 if smiles:211 desc += f" SMILES: {smiles[:80]}"212 return desc213 214 215def detect_pert_type(item, filename):216 """Detect pert_type from the input data or filename."""217 pert_type = item.get("pert_type", "")218 if not pert_type:219 pert_type = item.get("input", {}).get("pert_type", "")220 if not pert_type:221 if "chemical" in filename:222 pert_type = "chemical"223 elif "genetic" in filename:224 pert_type = "genetic"225 # Normalize: geneticRNv3 -> genetic226 if "genetic" in pert_type.lower():227 pert_type = "genetic"228 return pert_type229 230 231def construct_prediction_prompt(item, gene_summary_cache, pert_type):232 """Build the SUMMER-enhanced QA prompt for a 3way item."""233 inp = item["input"]234 pert_raw = inp["pert"]235 gene = inp["gene"]236 cell_type = inp["cell_type"]237 238 # Handle combo genetic perturbation (pert is a list)239 is_combo = isinstance(pert_raw, list)240 if is_combo:241 pert = "+".join(pert_raw) # e.g. "CBL+TGFBR2"242 else:243 pert = pert_raw244 245 # Gene summary (from KG - SUMMER's core enhancement)246 gene_desc = gene_summary_cache.get(247 gene, f"No knowledge graph information available for gene {gene}."248 )249 250 # Pert description251 if pert_type == "genetic":252 if is_combo:253 # Combine summaries for all genes in the combo254 parts = []255 for p in pert_raw:256 desc = gene_summary_cache.get(257 p, f"No knowledge graph information available for gene {p}."258 )259 parts.append(f"{p}: {desc}")260 pert_desc = "\n".join(parts)261 else:262 pert_desc = gene_summary_cache.get(263 pert, f"No knowledge graph information available for gene {pert}."264 )265 else:266 pert_desc = get_pert_description_chemical(item)267 268 # Reasoning paths from triplets269 triplets = item["output"].get("triplets_cell_conditioned",270 item["output"].get("triplets_all", []))271 reasoning_paths = format_reasoning_paths(triplets)272 273 # Basal expression context274 basal_context = item.get("basal_context", {})275 basal_text = format_basal_context(basal_context)276 277 # Cell description278 cell_desc = CELL_LINES.get(cell_type, f"{cell_type} cells.")279 280 # Select template281 if pert_type == "genetic":282 template = PROMPT_GENETIC_3WAY283 else:284 template = PROMPT_CHEMICAL_3WAY285 286 prompt = template.format(287 pert=pert,288 gene=gene,289 cell_short=cell_type,290 pert_desc=pert_desc,291 gene_desc=gene_desc,292 cell_desc=cell_desc,293 basal_text=basal_text,294 reasoning_paths=reasoning_paths,295 )296 return prompt297 298 299def parse_prediction(response):300 """Parse the 3-way prediction (up/down/unchanged) from LLM response."""301 # Primary: collect ALL <answer> tags and take the LAST one.302 # The model often echoes the template first (e.g. '<answer>up</answer> OR303 # <answer>down</answer> OR <answer>unchanged</answer>') before giving the real304 # answer, so re.search / the first match would be wrong.305 matches = re.findall(r'<answer>\s*(up|down|unchanged)\s*</answer>', response, re.IGNORECASE)306 if matches:307 return matches[-1].lower()308 309 # Fallback: look in last few lines310 response_lower = response.lower()311 last_lines = response_lower.split('\n')[-6:]312 last_text = ' '.join(last_lines)313 314 if "unchanged" in last_text:315 return "unchanged"316 317 has_down = bool(re.search(r'\bdown\b', last_text) and not re.search(r'knockdown', last_text)) or "decrease" in last_text318 has_up = bool(re.search(r'\bup\b', last_text)) or "increase" in last_text319 320 if has_down and not has_up:321 return "down"322 if has_up and not has_down:323 return "up"324 325 # Count-based fallback (full response)326 dec = response_lower.count("decrease") + response_lower.count("downregulat")327 inc = response_lower.count("increase") + response_lower.count("upregulat")328 unc = response_lower.count("unchanged") + response_lower.count("no significant")329 330 if unc > max(dec, inc):331 return "unchanged"332 if dec > inc:333 return "down"334 if inc > dec:335 return "up"336 return "unchanged" # conservative default337 338 339# ==============================================================================340# Main341# ==============================================================================342 343def main():344 parser = argparse.ArgumentParser(description="SUMMER Pipeline with vLLM (3way)")345 parser.add_argument("--kg_path", type=str, required=True,346 help="Path to knowledge graph directory")347 parser.add_argument("--input_dir", type=str, required=True,348 help="Directory containing input .jsonl files (3way dataset)")349 parser.add_argument("--output_dir", type=str, required=True,350 help="Directory to save output .jsonl files")351 parser.add_argument("--model_path", type=str, required=True,352 help="Path to the LLM model")353 parser.add_argument("--temperature", type=float, default=0.2)354 parser.add_argument("--max_tokens", type=int, default=4096)355 parser.add_argument("--tp_size", type=int, default=1)356 parser.add_argument("--gpu_memory_utilization", type=float, default=0.9)357 args = parser.parse_args()358 359 # 1. Initialize vLLM360 print(f"\n[Init] Loading vLLM model from {args.model_path}...", flush=True)361 llm = LLM(362 model=args.model_path,363 tensor_parallel_size=args.tp_size,364 trust_remote_code=True,365 dtype="bfloat16",366 max_model_len=8192,367 gpu_memory_utilization=args.gpu_memory_utilization,368 enforce_eager=True,369 )370 371 summary_sampling = SamplingParams(temperature=args.temperature, max_tokens=512)372 prediction_sampling = SamplingParams(temperature=args.temperature, max_tokens=args.max_tokens)373 374 # 2. Load Knowledge Graphs375 print("\n[Init] Loading Knowledge Graphs...", flush=True)376 kg_data = summer_pipeline.load_knowledge_graphs(args.kg_path)377 resources = summer_pipeline.build_kg_lookups(kg_data, max_degree=1000)378 379 # 3. Scan input files380 print("\n[Phase 0] Scanning input files...", flush=True)381 input_files = sorted(glob.glob(os.path.join(args.input_dir, "*.jsonl")))382 if not input_files:383 print(f"ERROR: No .jsonl files found in {args.input_dir}")384 return385 386 print(f" Found {len(input_files)} input files.", flush=True)387 388 # Collect all genes needing KG summaries389 all_genes = set()390 file_cache = {}391 392 for fpath in tqdm(input_files, desc="Reading files"):393 try:394 with open(fpath, 'r') as f:395 data = [json.loads(line) for line in f if line.strip()]396 file_cache[fpath] = data397 filename = os.path.basename(fpath)398 for item in data:399 inp = item["input"]400 all_genes.add(inp["gene"])401 pt = detect_pert_type(item, filename)402 if pt == "genetic":403 pert = inp["pert"]404 # pert may be a list for combo genetic perturbations405 if isinstance(pert, list):406 for p in pert:407 all_genes.add(p)408 else:409 all_genes.add(pert)410 except Exception as e:411 print(f" Error reading {fpath}: {e}")412 413 print(f" Total unique genes to summarize: {len(all_genes)}", flush=True)414 415 # 4. Summarize genes (Phase 1 - SUMMER's KG enhancement)416 print(f"\n[Phase 1] Generating KG summaries for {len(all_genes)} genes...", flush=True)417 gene_summary_cache = {}418 genes_to_process = []419 prompts_to_process = []420 421 for gene in sorted(all_genes):422 prompt = get_gene_summary_prompt(gene, resources)423 if prompt:424 genes_to_process.append(gene)425 prompts_to_process.append(prompt)426 else:427 gene_summary_cache[gene] = f"No knowledge graph information available for gene {gene}."428 429 if genes_to_process:430 chunk_size = 50431 print(f" Processing {len(genes_to_process)} genes in chunks of {chunk_size}...", flush=True)432 433 for i in range(0, len(genes_to_process), chunk_size):434 chunk_genes = genes_to_process[i:i + chunk_size]435 chunk_prompts = prompts_to_process[i:i + chunk_size]436 print(f" > Batch {i}-{i + len(chunk_genes)} / {len(genes_to_process)}", flush=True)437 438 try:439 outputs = llm.generate(chunk_prompts, summary_sampling, use_tqdm=False)440 for gene, output in zip(chunk_genes, outputs):441 if output.outputs:442 gene_summary_cache[gene] = output.outputs[0].text.strip()443 else:444 gene_summary_cache[gene] = ""445 except Exception as e:446 print(f" ! Error in batch {i}: {e}")447 for gene, prompt in zip(chunk_genes, chunk_prompts):448 try:449 out = llm.generate([prompt], summary_sampling, use_tqdm=False)450 gene_summary_cache[gene] = out[0].outputs[0].text.strip()451 except Exception:452 print(f" ! SKIPPING gene: {gene}")453 gene_summary_cache[gene] = f"Error processing gene {gene}."454 455 print(f" Cached summaries for {len(gene_summary_cache)} genes.", flush=True)456 457 # 5. Generate Predictions (Phase 2)458 print("\n[Phase 2] Generating Final Predictions...", flush=True)459 os.makedirs(args.output_dir, exist_ok=True)460 461 for fpath, qa_data in file_cache.items():462 filename = os.path.basename(fpath)463 output_path = os.path.join(args.output_dir, f"summer_pred_{filename}")464 465 if os.path.exists(output_path):466 print(f" Skipping {filename} (output already exists).")467 continue468 469 print(f"\n Processing {filename} ({len(qa_data)} items)...", flush=True)470 471 if not qa_data:472 continue473 474 pert_type = detect_pert_type(qa_data[0], filename)475 print(f" pert_type={pert_type}", flush=True)476 477 # Build prompts478 prompts = []479 valid_items = []480 for item in qa_data:481 try:482 p = construct_prediction_prompt(item, gene_summary_cache, pert_type)483 if len(p) > 32000:484 print(f" Warning: Skipping item {item.get('id')} (prompt too long).")485 continue486 prompts.append(p)487 valid_items.append(item)488 except Exception as e:489 print(f" Skipping item {item.get('id')}: {e}")490 491 if not prompts:492 print(f" No valid prompts for {filename}.")493 continue494 495 # Run inference in chunks496 chunk_size = 200497 all_outputs = []498 for ci in range(0, len(prompts), chunk_size):499 chunk_prompts = prompts[ci:ci + chunk_size]500 print(f" Inference chunk {ci}-{ci + len(chunk_prompts)} / {len(prompts)}", flush=True)501 try:502 chunk_outputs = llm.generate(chunk_prompts, prediction_sampling, use_tqdm=True)503 all_outputs.extend(chunk_outputs)504 except Exception as e:505 print(f" ! Error in chunk {ci}: {e}")506 for p in chunk_prompts:507 try:508 out = llm.generate([p], prediction_sampling, use_tqdm=False)509 all_outputs.extend(out)510 except Exception:511 all_outputs.append(None)512 513 # Write output (matching Qwen's 4-key format)514 with open(output_path, "w", encoding="utf-8") as f_out:515 for item, summer_prompt, output in zip(valid_items, prompts, all_outputs):516 if output is None or not output.outputs:517 model_output = "ERROR: Failed to generate output."518 else:519 model_output = output.outputs[0].text520 521 ground_truth = item.get("response")522 523 result = {524 "source_file": filename,525 "prompt": summer_prompt,526 "ground_truth_response": ground_truth,527 "model_output": model_output,528 }529 f_out.write(json.dumps(result, ensure_ascii=False) + "\n")530 531 print(f" Wrote {len(valid_items)} predictions to {output_path}", flush=True)532 533 print("\n[Done] SUMMER pipeline complete.", flush=True)534 535 536if __name__ == "__main__":537 main()538