#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ Clinical Trial Matching Pipeline - Gradio Web Interface This interface allows users to: 1. Configure models (tagger, embedder, LLM) 2. Upload trial space database OR load pre-embedded trials 3. Upload patient notes or enter patient summary 4. Get ranked trial recommendations with eligibility predictions """ import gradio as gr import pandas as pd import numpy as np import torch import re import os import json from typing import List, Tuple, Optional, Dict from pathlib import Path import tempfile # HuggingFace imports from transformers import ( AutoTokenizer, AutoModelForSequenceClassification, pipeline ) from sentence_transformers import SentenceTransformer from datasets import load_dataset # Try to import configuration try: import config HAS_CONFIG = True print("āœ“ Found config.py - will auto-load models on startup") except ImportError: HAS_CONFIG = False print("ā—‹ No config.py found - using manual model loading") # Global state to hold loaded models and embedded trials class AppState: def __init__(self): self.tagger_model = None self.tagger_tokenizer = None self.embedder_model = None self.embedder_tokenizer = None self.llm_model = None self.llm_tokenizer = None self.trial_checker_model = None self.trial_checker_tokenizer = None self.boilerplate_checker_model = None self.boilerplate_checker_tokenizer = None self.trial_spaces_df = None self.trial_embeddings = None self.trial_preview_df = None self.device = "cuda" if torch.cuda.is_available() else "cpu" # Store auto-load status messages to display in UI self.auto_load_status = { "tagger": "", "embedder": "", "llm": "", "trial_checker": "", "boilerplate_checker": "", "trials": "" } def reset_trials(self): self.trial_spaces_df = None self.trial_embeddings = None self.trial_preview_df = None state = AppState() # ============================================================================ # UTILITY FUNCTIONS # ============================================================================ def split_into_excerpts(text: str) -> List[str]: """Split text into sentence-level excerpts.""" if not text or pd.isna(text): return [] t = re.sub(r'[\n\r]+', ' ', text.strip()) t = re.sub(r'\s+', ' ', t) if not t: return [] t2 = t.replace(". ", "") parts = [p.strip() for p in t2.split("") if p.strip()] return parts def truncate_text(text: str, tokenizer, max_tokens: int = 1500) -> str: """Truncate text to a maximum number of tokens.""" return tokenizer.decode( tokenizer.encode(text, add_special_tokens=True, truncation=True, max_length=max_tokens), skip_special_tokens=True ) def format_probability_visual(val, is_exclusion=False): """ Helper to format probabilities with visual indicators (emojis) for the dataframe. """ try: val_float = float(val) except: return val # Logic for Eligibility (High is good) if not is_exclusion: if val_float >= 0.8: return f"šŸŸ¢ **{val_float:.2f}**" elif val_float >= 0.5: return f"šŸŸ” {val_float:.2f}" else: return f"šŸ”“ {val_float:.2f}" # Logic for Exclusion (High is bad) else: if val_float >= 0.5: return f"šŸ”“ **{val_float:.2f}**" # High exclusion prob is bad elif val_float >= 0.2: return f"šŸŸ” {val_float:.2f}" else: return f"šŸŸ¢ {val_float:.2f}" # Low exclusion prob is good # ============================================================================ # TRIAL SPACE EXTRACTION CONSTANTS # ============================================================================ MAX_EMBEDDER_SEQ_LEN = 2500 MAX_LONGTEXT_SEQ_LEN = 110000 MAX_TRIAL_CHECKER_LENGTH = 4096 MAX_BOILERPLATE_CHECKER_LENGTH = 3192 REASONING_MARKER = "assistantfinal" BOILERPLATE_MARKER = "Boilerplate" TRIAL_SPACE_PROMPT_HEADER = ( "You are an expert clinical oncologist with an encyclopedic knowledge of cancer and its treatments.\n" "Your job is to review a clinical trial document and extract a list of structured clinical spaces that are eligible for that trial.\n" "A clinical space is defined as a unique combination of patient age range, sex (if any sex criteria), cancer primary site, histology, which treatments a patient must have received, " "which treatments a patient must not have received, cancer burden (eg presence of metastatic disease; this also includes cancer type-specific prognostic scores, risk indices, or categories), tumor biomarkers (such as " "germline or somatic gene mutations or alterations, or protein expression on tumor), that a patient must have or must not have to " "be eligible for the trial. \n" "With respect to sex criteria: For cancers originating in organs only present in one sex, you must assume the sex criteria even if not stated explicitly.\n" "For example, a trial space for uterine, ovarian, vulvar, vaginal, or fallopian tube cancer must be assumed to be for female patients.\n" "Similarly, a trial space for testicular, penile, or prostate cancer must be assumed to be for male patients.\n" "For all other cancer types (including breast cancer), you shoulud assume the trial is open to both sexes unless the clinical trial document states otherwise.\n" "Trials often specify that a particular treatment is excluded only if it was given within a short period of time, for example 14 days, " "one month, etc , prior to trial start. This is called a washout period. Do not include this type of time-specific treatment washout " "eligibility criteria in your output at all.\n" "Some trials have only one space, while others have several. Do not output a space that contains multiple cancer types and/or histologies. " "Instead, generate separate spaces for each cancer type/histology combination.\n" "CRITICAL: Each trial space must contain all information necessary to define that space on its own. It may not refer to other previously " "defined spaces for the same trial, since for later use, the spaces will be extracted and separated from each other. YOU MAY NOT include " "text describing a given space that refers to a previous space; eg, \"Same as above\"-style output is not allowed!\n" "For biomarkers, if the trial specifies whether the biomarker will be assessed during screening, note that.\n" "Spell out cancer types; do not abbreviate them. For example, write \"non-small cell lung cancer\" rather than \"NSCLC\".\n" "Structure your output like this, as a list of spaces, with spaces separated by newlines, as below. STRICTLY adhere to the formatting.\n" "1. Age range allowed: . Sex allowed: . Cancer type allowed: . Histology allowed: . Cancer burden allowed: . Prior treatment required: . Prior treatment excluded: . Biomarkers required: . Biomarkers excluded: . \n" "2. Cancer type allowed: , etc.\n" "If a concept is not relevant, such as if there are no prior treatents required, simply output NA for that concept.\n" "CRITICAL: Anytime you provide a list for a particular concept, you must be completely clear on whether \"or\" versus \"and\" logic applies " "to the list. For example, do not output \"EGFR L858R mutant, TP53 mutant\"; if both are required, output \"EGFR L858R mutant and TP53 mutant\". " "As another example, do not output \"ER+, PR+\"; if the patient can have either an ER or a PR positive tumor, output \"ER+ or PR+\".\n" "NEVER put a newline within a single trial space.\n" "After you output the trial spaces, output a newline, then the text \"Boilerplate exclusions:\" VERBATIM, then another newline.\n" "Then, list exclusion criteria described in the trial text that are unrelated to the trial space definitions. Such exclusions tend to be common " "to clinical trials in general.\n" "Common boilerplate exclusion criteria include a history of pneumonitis, heart failure, renal dysfunction, liver dysfunction, uncontrolled brain " "metastases, HIV or hepatitis, and poor performance status.\n" "ALWAYS output plain text only. NEVER output unicode, Markdown, or tables.\n" ) TRIAL_SPACE_PROMPT_SUFFIX = ( "Now, generate your list of the trial space(s), followed by any boilerplate exclusions, formatted as above.\n" "Do not provide any introductory, explanatory, concluding, or disclaimer text.\n" "Reminder: Treatment history is an important component of trial space definitions, but treatment history \"washout\" requirements that are " "described as applying only in a given period of time prior to trial treatment MUST BE IGNORED.\n" "CRITICAL: A given trial space MUST NEVER refer to another previously defined space. You must NEVER output text like \"same as #1\" or " "\"same criteria as above.\" Instead, you MUST REPEAT all relevant criteria for each new space SO THAT IT STANDS ON ITS OWN. A user who later " "looks at the text for one space will not have access to text for other spaces, and so output like \"Same criteria as #1...\" renders a space useless!" ) # ============================================================================ # AUTO-LOADING FROM CONFIG # ============================================================================ def auto_load_models_from_config(): """Auto-load models specified in config.py""" if not HAS_CONFIG: return print("\n" + "="*70) print("AUTO-LOADING MODELS FROM CONFIG") print("="*70) # Load tagger if config.MODEL_CONFIG.get("tagger"): print(f"\n[1/5] Loading tagger: {config.MODEL_CONFIG['tagger']}") status, _ = load_tagger_model(config.MODEL_CONFIG["tagger"]) state.auto_load_status["tagger"] = status print(status) # Load embedder if config.MODEL_CONFIG.get("embedder"): print(f"\n[2/5] Loading embedder: {config.MODEL_CONFIG['embedder']}") status, _, _ = load_embedder_model(config.MODEL_CONFIG["embedder"]) state.auto_load_status["embedder"] = status print(status) # Load LLM if config.MODEL_CONFIG.get("llm"): print(f"\n[3/5] Loading LLM: {config.MODEL_CONFIG['llm']}") status, _ = load_llm_model(config.MODEL_CONFIG["llm"]) state.auto_load_status["llm"] = status print(status) # Load trial checker if config.MODEL_CONFIG.get("trial_checker"): print(f"\n[4/5] Loading trial checker: {config.MODEL_CONFIG['trial_checker']}") status, _ = load_trial_checker(config.MODEL_CONFIG["trial_checker"]) state.auto_load_status["trial_checker"] = status print(status) # Load boilerplate checker if config.MODEL_CONFIG.get("boilerplate_checker"): print(f"\n[5/5] Loading boilerplate checker: {config.MODEL_CONFIG['boilerplate_checker']}") status, _ = load_boilerplate_checker(config.MODEL_CONFIG["boilerplate_checker"]) state.auto_load_status["boilerplate_checker"] = status print(status) print("\n" + "="*70) print("MODEL AUTO-LOADING COMPLETE") print("="*70 + "\n") def auto_load_trials_from_config(): """Auto-load trial database from config.py - prefers pre-embedded over fresh embedding.""" if not HAS_CONFIG: return # Check for pre-embedded trials first (much faster) if hasattr(config, 'PREEMBEDDED_TRIALS') and config.PREEMBEDDED_TRIALS: preembed_path = config.PREEMBEDDED_TRIALS print("\n" + "="*70) print(f"AUTO-LOADING PRE-EMBEDDED TRIALS: {preembed_path}") print("="*70) status, preview = load_preembedded_trials(preembed_path) state.auto_load_status["trials"] = status # Store the preview so it can be displayed in the UI state.trial_preview_df = preview print("="*70) print("PRE-EMBEDDED TRIALS AUTO-LOADING COMPLETE") print("="*70 + "\n") return # Fall back to fresh embedding if no pre-embedded trials specified if not hasattr(config, 'DEFAULT_TRIAL_DB') or not config.DEFAULT_TRIAL_DB: print("ā—‹ No trial database specified in config") return if not os.path.exists(config.DEFAULT_TRIAL_DB): print(f"āœ— Default trial database not found: {config.DEFAULT_TRIAL_DB}") state.auto_load_status["trials"] = f"āœ— Trial database file not found: {config.DEFAULT_TRIAL_DB}" return if state.embedder_model is None: print("ā—‹ Embedder not loaded yet - skipping trial database auto-load") state.auto_load_status["trials"] = "ā—‹ Waiting for embedder model to be loaded..." return print("\n" + "="*70) print(f"AUTO-LOADING TRIAL DATABASE: {config.DEFAULT_TRIAL_DB}") print("="*70) # Create a temporary file-like object class FilePath: def __init__(self, path): self.name = path status, preview = load_and_embed_trials(FilePath(config.DEFAULT_TRIAL_DB), show_progress=True) state.auto_load_status["trials"] = status # Store the preview so it can be displayed in the UI state.trial_preview_df = preview print("="*70) print("TRIAL DATABASE AUTO-LOADING COMPLETE") print("="*70 + "\n") # ============================================================================ # MODEL LOADING FUNCTIONS # ============================================================================ def load_tagger_model(model_path: str) -> Tuple[str, str]: """Load TinyBERT tagger model.""" try: state.tagger_tokenizer = AutoTokenizer.from_pretrained(model_path) state.tagger_model = pipeline( "text-classification", model=model_path, tokenizer=state.tagger_tokenizer, device=0 if state.device == "cuda" else -1, truncation=True, padding="max_length", max_length=128 ) return f"āœ“ Tagger model loaded from {model_path}", "" except Exception as e: return f"āœ— Error loading tagger model: {str(e)}", str(e) def load_embedder_model(model_path: str) -> Tuple[str, str, str]: """Load sentence transformer embedder model.""" try: # Check if trials are already loaded will_need_reembed = state.trial_spaces_df is not None and len(state.trial_spaces_df) > 0 if will_need_reembed: warning_msg = f"\nāš ļø Warning: {len(state.trial_spaces_df)} trials are currently loaded. They will need to be re-embedded with the new model." else: warning_msg = "" state.embedder_model = SentenceTransformer(model_path, device=state.device, trust_remote_code=True) state.embedder_tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) # Set the instruction prompt try: state.embedder_model.prompts['query'] = ( "Instruct: Given a cancer patient summary, retrieve clinical trial options " "that are reasonable for that patient; or, given a clinical trial option, " "retrieve cancer patients who are reasonable candidates for that trial." ) except: pass try: state.embedder_model.max_seq_length = MAX_EMBEDDER_SEQ_LEN except: pass success_msg = f"āœ“ Embedder model loaded from {model_path}{warning_msg}" # If trials were loaded, invalidate embeddings if will_need_reembed: state.trial_embeddings = None success_msg += "\nā†’ Trial embeddings cleared. Please reload trial database to re-embed." return success_msg, "", warning_msg except Exception as e: return f"āœ— Error loading embedder model: {str(e)}", str(e), "" def load_llm_model(model_path: str) -> Tuple[str, str]: """Load LLM for patient summarization.""" try: # Check if vLLM is available try: from vllm import LLM, SamplingParams # Determine tensor parallel size gpu_count = torch.cuda.device_count() tp_size = min(gpu_count, 4) if gpu_count > 1 else 1 state.llm_model = LLM( model=model_path, tensor_parallel_size=tp_size, gpu_memory_utilization=0.60, max_model_len=15000 ) state.llm_tokenizer = state.llm_model.get_tokenizer() return f"āœ“ LLM loaded from {model_path} (vLLM, tp={tp_size})", "" except ImportError: # Fallback to HuggingFace transformers from transformers import AutoModelForCausalLM state.llm_tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) state.llm_model = AutoModelForCausalLM.from_pretrained( model_path, torch_dtype=torch.float16 if state.device == "cuda" else torch.float32, device_map="auto", trust_remote_code=True ) return f"āœ“ LLM loaded from {model_path} (HuggingFace)", "" except Exception as e: return f"āœ— Error loading LLM: {str(e)}", str(e) def load_trial_checker(model_path: str) -> Tuple[str, str]: """Load ModernBERT trial checker.""" try: state.trial_checker_tokenizer = AutoTokenizer.from_pretrained(model_path) state.trial_checker_model = AutoModelForSequenceClassification.from_pretrained( model_path, torch_dtype=torch.float16 if state.device == "cuda" else torch.float32 ).to(state.device) state.trial_checker_model.eval() return f"āœ“ Trial checker loaded from {model_path}", "" except Exception as e: return f"āœ— Error loading trial checker: {str(e)}", str(e) def load_boilerplate_checker(model_path: str) -> Tuple[str, str]: """Load ModernBERT boilerplate checker.""" try: state.boilerplate_checker_tokenizer = AutoTokenizer.from_pretrained(model_path) state.boilerplate_checker_model = AutoModelForSequenceClassification.from_pretrained( model_path, torch_dtype=torch.float16 if state.device == "cuda" else torch.float32 ).to(state.device) state.boilerplate_checker_model.eval() return f"āœ“ Boilerplate checker loaded from {model_path}", "" except Exception as e: return f"āœ— Error loading boilerplate checker: {str(e)}", str(e) # ============================================================================ # TRIAL SPACE PROCESSING (WITH PRE-EMBEDDING SUPPORT) # ============================================================================ def load_preembedded_trials(path_or_url: str) -> Tuple[str, pd.DataFrame]: """Load pre-embedded trial database from a local parquet file or a Huggingface URL.""" try: print(f"\n{'='*70}") print(f"LOADING PRE-EMBEDDED TRIALS") print(f"{'='*70}") print(f"Loading from: {path_or_url}") # Check if it's a URL or a local path if path_or_url.startswith("http"): # It's a URL, load from Huggingface print("Detected URL, loading from Huggingface Hub...") dataset = load_dataset("parquet", data_files=path_or_url, split='train') df = dataset.to_pandas() print(f"āœ“ Loaded {len(df)} trials from Hub") else: # It's a local path parquet_path = path_or_url if not parquet_path.endswith('.parquet'): parquet_path = parquet_path + '.parquet' # Check file exists if not os.path.exists(parquet_path): return f"āœ— Pre-embedded parquet file not found: {parquet_path}", None # Load parquet file print(f"Loading trial dataframe with embeddings...") df = pd.read_parquet(parquet_path) print(f"āœ“ Loaded {len(df)} trials from local file") # Check for embedding column if 'embedding' not in df.columns: return f"āœ— Parquet file missing 'embedding' column: {path_or_url}", None # Extract embeddings from the column and convert to numpy array print(f"Extracting embeddings...") embeddings = np.array(df['embedding'].tolist()) print(f"āœ“ Extracted embeddings: {embeddings.shape}") # Remove embedding column from dataframe (not needed in the df itself) df_without_embeddings = df.drop(columns=['embedding']) # Store in state state.trial_spaces_df = df_without_embeddings state.trial_embeddings = embeddings print(f"{'='*70}") print(f"PRE-EMBEDDED TRIALS LOADED SUCCESSFULLY") print(f"{'='*70}\n") preview = df_without_embeddings[['nct_id', 'this_space']].head(10) return f"āœ“ Loaded {len(df)} pre-embedded trials from {path_or_url}", preview except Exception as e: import traceback traceback.print_exc() return f"āœ— Error loading pre-embedded trials: {str(e)}", None def load_and_embed_trials(file, show_progress: bool = False) -> Tuple[str, pd.DataFrame]: """Load trial spaces CSV/Excel and embed them.""" try: if state.embedder_model is None: return "āœ— Please load the embedder model first!", None # Read file if file.name.endswith('.csv'): df = pd.read_csv(file.name) elif file.name.endswith(('.xlsx', '.xls')): df = pd.read_excel(file.name) else: return "āœ— Unsupported file format. Use CSV or Excel.", None # Check required columns required_cols = ['nct_id', 'this_space', 'trial_text', 'trial_boilerplate_text'] missing = [col for col in required_cols if col not in df.columns] if missing: return f"āœ— Missing required columns: {', '.join(missing)}", None # Clean data df = df[~df['this_space'].isnull()].copy() df['trial_boilerplate_text'] = df['trial_boilerplate_text'].fillna('') # Prepare texts for embedding df['this_space_trunc'] = df['this_space'].apply( lambda x: truncate_text(str(x), state.embedder_tokenizer, max_tokens=MAX_EMBEDDER_SEQ_LEN) ) # Add instruction prefix prefix = ( "Instruct: Given a cancer patient summary, retrieve clinical trial options " "that are reasonable for that patient; or, given a clinical trial option, " "retrieve cancer patients who are reasonable candidates for that trial. " ) texts_to_embed = [prefix + txt for txt in df['this_space_trunc'].tolist()] # Embed with progress if not show_progress: gr.Info(f"Embedding {len(df)} trial spaces...") else: print(f"Embedding {len(df)} trial spaces...") with torch.no_grad(): embeddings = state.embedder_model.encode( texts_to_embed, batch_size=64, convert_to_tensor=True, normalize_embeddings=True, show_progress_bar=show_progress, prompt='query' ) # Store in state state.trial_spaces_df = df state.trial_embeddings = embeddings.cpu().numpy() preview = df[['nct_id', 'this_space']].head(10) success_msg = f"āœ“ Loaded and embedded {len(df)} trial spaces" if show_progress: print(success_msg) return success_msg, preview except Exception as e: return f"āœ— Error processing trials: {str(e)}", None # ============================================================================ # PATIENT NOTE PROCESSING # ============================================================================ def process_patient_notes(file, prob_threshold: float = 0.1) -> Tuple[str, str]: """Process patient notes through tagger and create long note.""" try: if state.tagger_model is None: return "āœ— Please load the tagger model first!", "" # Read file if file.name.endswith('.csv'): df = pd.read_csv(file.name) elif file.name.endswith(('.xlsx', '.xls')): df = pd.read_excel(file.name) else: return "āœ— Unsupported file format. Use CSV or Excel.", "" # Check required columns if 'date' not in df.columns or 'text' not in df.columns: return "āœ— File must contain 'date' and 'text' columns", "" # Sort by date df['date'] = pd.to_datetime(df['date'], errors='coerce') df = df.sort_values('date').reset_index(drop=True) # Extract all excerpts all_excerpts = [] all_dates = [] all_note_types = [] for idx, row in df.iterrows(): excerpts = split_into_excerpts(str(row['text'])) note_type = row.get('note_type', 'clinical_note') for exc in excerpts: all_excerpts.append(exc) all_dates.append(row['date']) all_note_types.append(note_type) if not all_excerpts: return "āœ— No valid excerpts extracted from notes", "" gr.Info(f"Tagging {len(all_excerpts)} excerpts...") # Run tagger predictions = state.tagger_model(all_excerpts, batch_size=256) # Extract positive excerpts excerpts_df = pd.DataFrame({ 'excerpt': all_excerpts, 'date': all_dates, 'note_type': all_note_types, 'label': [p['label'] for p in predictions], 'score': [p['score'] for p in predictions] }) # Calculate positive probability excerpts_df['positive_prob'] = np.where( excerpts_df['label'] == 'NEGATIVE', 1.0 - excerpts_df['score'], excerpts_df['score'] ) # Filter by threshold keep = excerpts_df[excerpts_df['positive_prob'] > prob_threshold].copy() # FIX: Capture the raw count of excerpts *before* grouping raw_keep_count = len(keep) if len(keep) == 0: return "āœ— No excerpts passed the threshold", "" # Group by date and note type keep['date_str'] = keep['date'].dt.strftime('%Y-%m-%d') keep = keep.groupby(['date_str', 'note_type'])['excerpt'].agg(lambda x: ' '.join(x)).reset_index() keep['date_text'] = ( keep['date_str'] + " " + keep['note_type'] + " " + keep['excerpt'] ) # Create long note long_note = "\n".join(keep['date_text'].tolist()) # FIX: Display the raw count in the stats message stats = ( f"Processed {len(df)} notes ā†’ {len(all_excerpts)} excerpts ā†’ " f"{raw_keep_count} relevant excerpts (threshold={prob_threshold})" ) return stats, long_note except Exception as e: return f"āœ— Error processing notes: {str(e)}", "" def summarize_patient_history(long_note: str) -> Tuple[str, str]: """Summarize patient long note using LLM and split into summary and boilerplate sections.""" try: if state.llm_model is None: return "āœ— Please load the LLM model first!", "" if not long_note or len(long_note.strip()) == 0: return "āœ— No patient history to summarize", "" # Truncate if needed tokens = state.llm_tokenizer.encode(long_note, add_special_tokens=False) max_tokens = MAX_LONGTEXT_SEQ_LEN # Leave room for prompt and response if len(tokens) > max_tokens: half = max_tokens // 2 first_part = state.llm_tokenizer.decode(tokens[:half]) last_part = state.llm_tokenizer.decode(tokens[-half:]) patient_text = first_part + " ... " + last_part else: patient_text = long_note # Build prompt messages = [ {'role': 'system', 'content': 'Reasoning: high'}, {'role': 'user', 'content': """You are an experienced clinical oncology history summarization bot. Your job is to construct a summary of the cancer history for a patient based on an excerpt of the patient's electronic health record. The text in the excerpt is provided in chronological order. Each paragraph in the excerpt represents a summary of a clinical document written on the date indicated in the paragraph. Document the patient's most recent age; sex; cancer type/primary site (eg breast cancer, lung cancer, etc); histology (eg adenocarcinoma, squamous carcinoma, etc); current extent (localized, advanced, metastatic, etc); biomarkers (genomic results, protein expression, etc); and treatment history (surgery, radiation, chemotherapy/targeted therapy/immunotherapy, etc, including start and stop dates and best response if known). Do not consider localized basal cell or squamous carcinomas of the skin, or colon polyps, to be cancers for your purposes. Do not include the patient's name, but do include relevant dates whenever documented, including dates of diagnosis and start/stop dates of each treatment. If a patient has a history of more than one cancer, document the cancers one at a time. CRITICAL: Format your response as free text ONLY. Do NOT output markdown, Unicode, or tables. Also document any history of conditions that might meet "boilerplate" exclusion criteria, including uncontrolled brain metastases, lack of measurable disease, congestive heart failure, pneumonitis, renal dysfunction, liver dysfunction, and HIV or hepatitis infection. For each of these, present the evidence from the history that the patient has a history of such a condition, including dates. Clearly separate the "boilerplate" section by labeling it "Boilerplate: " before describing any such conditions. Here is an example of the desired output format: Age: 70 Sex: Male Cancer type: Lung cancer Histology: Adenocarcinoma Current extent: Metastatic Biomarkers: PD-L1 75%, KRAS G12C mutant Treatment history: # 1/5/2020-2/5/2021: carboplatin/pemetrexed/pembrolizumab # 1/2021: Palliative radiation to progressive spinal metastases # 3/2021-present: docetaxel Boilerplate: No evidence of common boilerplate exclusion criteria """ + "The excerpt for you to summarize is:\n" + patient_text + """\nNow, write your summary. Do not add preceding text before the abstraction, and do not add notes or commentary afterwards. This will not be used for clinical care, so do not write any disclaimers or cautionary notes."""} ] gr.Info("Summarizing patient history with LLM...") # Check if using vLLM or HuggingFace if hasattr(state.llm_model, 'generate') and hasattr(state.llm_model, 'get_tokenizer'): # vLLM from vllm import SamplingParams prompt = state.llm_tokenizer.apply_chat_template( conversation=messages, add_generation_prompt=True, tokenize=False ) response = state.llm_model.generate( [prompt], SamplingParams( temperature=0.0, top_k=1, max_tokens=7500, repetition_penalty=1.2 ) ) output = response[0].outputs[0].text else: # HuggingFace input_ids = state.llm_tokenizer.apply_chat_template( conversation=messages, add_generation_prompt=True, return_tensors="pt" ).to(state.device) with torch.no_grad(): outputs = state.llm_model.generate( input_ids, max_new_tokens=7500, temperature=0.00, do_sample=True, repetition_penalty=1.2 ) output = state.llm_tokenizer.decode(outputs[0], skip_special_tokens=True) if REASONING_MARKER in output: cleaned = output.split(REASONING_MARKER, 1)[-1] else: cleaned = output # 2. Handle Boilerplate Marker (Updated to split by line) lines = cleaned.splitlines(keepends=True) marker_line_index = -1 # Find which line contains the marker for i, line in enumerate(lines): if BOILERPLATE_MARKER in line: marker_line_index = i break if marker_line_index != -1: # Join lines before the marker line pre = "".join(lines[:marker_line_index]) # Join lines after the marker line (skipping the marker line itself) post = "".join(lines[marker_line_index + 1:]) summary_text = pre boilerplate_text = post else: # If marker not found, the full text goes to both summary_text = cleaned boilerplate_text = cleaned return summary_text, boilerplate_text except Exception as e: return f"āœ— Error summarizing: {str(e)}", "" # ============================================================================ # TRIAL SPACE EXTRACTION # ============================================================================ def extract_trial_spaces(trial_text: str) -> str: """Extract trial spaces and boilerplate criteria from trial text using LLM.""" try: if state.llm_model is None: return "āœ— Please load the LLM model first!" if not trial_text or len(trial_text.strip()) == 0: return "āœ— No trial text provided" # Build prompt messages messages = [ {"role": "system", "content": "Reasoning: high."}, { "role": "user", "content": ( TRIAL_SPACE_PROMPT_HEADER + "Here is a clinical trial document:\n" + str(trial_text) + "\n" + TRIAL_SPACE_PROMPT_SUFFIX ), }, ] gr.Info("Extracting trial spaces with LLM...") # Check if using vLLM or HuggingFace if hasattr(state.llm_model, 'generate') and hasattr(state.llm_model, 'get_tokenizer'): # vLLM from vllm import SamplingParams prompt = state.llm_tokenizer.apply_chat_template( conversation=messages, add_generation_prompt=True, tokenize=False ) response = state.llm_model.generate( [prompt], SamplingParams( temperature=0.0, top_k=1, max_tokens=7500, repetition_penalty=1.3 ) ) output = response[0].outputs[0].text else: # HuggingFace input_ids = state.llm_tokenizer.apply_chat_template( conversation=messages, add_generation_prompt=True, return_tensors="pt" ).to(state.device) with torch.no_grad(): outputs = state.llm_model.generate( input_ids, max_new_tokens=7500, temperature=0.0, do_sample=False, repetition_penalty=1.3 ) output = state.llm_tokenizer.decode(outputs[0], skip_special_tokens=True) # Extract just the assistant response if REASONING_MARKER in output: output = output.split(REASONING_MARKER)[-1] # Clean up reasoning markers if present if REASONING_MARKER in output: output = output.split(REASONING_MARKER, 1)[-1] output = output.strip() return output except Exception as e: return f"āœ— Error extracting trial spaces: {str(e)}" # ============================================================================ # TRIAL MATCHING # ============================================================================ def match_trials(patient_summary: str, patient_boilerplate: str, top_k: int = 20) -> pd.DataFrame: """Match patient to trials and run checkers.""" try: if state.embedder_model is None: raise ValueError("Embedder model not loaded") if state.trial_embeddings is None: raise ValueError("Trial spaces not loaded") if state.trial_checker_model is None: raise ValueError("Trial checker model not loaded") if state.boilerplate_checker_model is None: raise ValueError("Boilerplate checker model not loaded") # Embed patient summary prefix = ( "Instruct: Given a cancer patient summary, retrieve clinical trial options " "that are reasonable for that patient; or, given a clinical trial option, " "retrieve cancer patients who are reasonable candidates for that trial. " ) patient_text = truncate_text(patient_summary, state.embedder_tokenizer, max_tokens=MAX_EMBEDDER_SEQ_LEN) patient_text_with_prefix = prefix + patient_text gr.Info("Embedding patient summary...") with torch.no_grad(): patient_emb = state.embedder_model.encode( [patient_text_with_prefix], convert_to_tensor=True, normalize_embeddings=True, prompt='query' ) # Calculate similarities patient_emb_np = patient_emb.cpu().numpy() similarities = np.dot(state.trial_embeddings, patient_emb_np.T).squeeze() # Get top-k top_indices = np.argsort(similarities)[::-1][:top_k] # Get top trials top_trials = state.trial_spaces_df.iloc[top_indices].copy() top_trials['similarity_score'] = similarities[top_indices] gr.Info(f"Running eligibility checks on top {len(top_trials)} trials...") # Run trial checker trial_check_inputs = [ f"{row['this_space']}\nNow here is the patient summary:{patient_summary}" for _, row in top_trials.iterrows() ] trial_check_encodings = state.trial_checker_tokenizer( trial_check_inputs, truncation=True, max_length=MAX_TRIAL_CHECKER_LENGTH, padding=True, return_tensors='pt' ).to(state.device) with torch.no_grad(): trial_check_outputs = state.trial_checker_model(**trial_check_encodings) trial_probs = torch.softmax(trial_check_outputs.logits, dim=1)[:, 1].cpu().numpy() top_trials['eligibility_probability'] = trial_probs # Run boilerplate checker boilerplate_check_inputs = [ f"Patient history: {patient_boilerplate}\nTrial exclusions:{row['trial_boilerplate_text']}" for _, row in top_trials.iterrows() ] boilerplate_check_encodings = state.boilerplate_checker_tokenizer( boilerplate_check_inputs, truncation=True, max_length=MAX_BOILERPLATE_CHECKER_LENGTH, padding=True, return_tensors='pt' ).to(state.device) with torch.no_grad(): boilerplate_check_outputs = state.boilerplate_checker_model(**boilerplate_check_encodings) boilerplate_probs = torch.softmax(boilerplate_check_outputs.logits, dim=1)[:, 1].cpu().numpy() top_trials['exclusion_probability'] = boilerplate_probs # Sort by eligibility probability top_trials = top_trials.sort_values('eligibility_probability', ascending=False) # Apply visual formatting for the display table top_trials['eligibility_display'] = top_trials['eligibility_probability'].apply( lambda x: format_probability_visual(x, is_exclusion=False) ) top_trials['exclusion_display'] = top_trials['exclusion_probability'].apply( lambda x: format_probability_visual(x, is_exclusion=True) ) top_trials['similarity_display'] = top_trials['similarity_score'].apply( lambda x: f"{x:.3f}" ) # Select columns for display - use the Display versions for the UI display_cols = [ 'nct_id', 'eligibility_display', 'exclusion_display', 'similarity_display', 'this_space' ] result_df = top_trials[display_cols].reset_index(drop=True) # Rename columns for better UI reading result_df.columns = [ 'NCT ID', 'Eligibility', 'Exclusion', 'Similarity', 'Criteria Space' ] return result_df except Exception as e: gr.Error(f"Error matching trials: {str(e)}") return pd.DataFrame() def get_trial_details(df: pd.DataFrame, evt: gr.SelectData) -> str: """Get full trial details when user clicks on a row.""" try: if df is None or len(df) == 0: return "No trial selected" row_idx = evt.index[0] # Map renamed columns back to logic nct_id = df.iloc[row_idx]['NCT ID'] this_space = df.iloc[row_idx]['Criteria Space'] # Find the specific trial space in original dataframe # Match both NCT ID and the exact trial space text matching_rows = state.trial_spaces_df[ (state.trial_spaces_df['nct_id'] == nct_id) & (state.trial_spaces_df['this_space'] == this_space) ] if len(matching_rows) == 0: return f"Error: Could not find matching trial space for {nct_id}" trial_row = matching_rows.iloc[0] # Create clinicaltrials.gov link ct_gov_link = f"https://clinicaltrials.gov/study/{nct_id}" details = f""" # Trial Details: {nct_id} **šŸ”— [View on ClinicalTrials.gov]({ct_gov_link})** --- ## Eligibility Criteria Summary (Selected Space) {trial_row['this_space']} ## Full Trial Text {trial_row['trial_text']} ## Boilerplate Exclusions {trial_row['trial_boilerplate_text']} """ return details except Exception as e: return f"Error retrieving trial details: {str(e)}" # ============================================================================ # GRADIO INTERFACE # ============================================================================ def create_interface(): # Attempt to load sample patient notes for pre-population sample_notes_path = None try: gr.Info("Loading sample patient notes...") # Correct raw URL for the CSV file sample_url = "https://huggingface.co/datasets/ksg-dfci/mmai-synthetic/raw/main/sample_patient_notes.csv" sample_df = pd.read_csv(sample_url) # Create a temporary file to hold the sample notes with tempfile.NamedTemporaryFile(delete=False, mode='w', suffix='.csv', newline='') as tmpfile: sample_df.to_csv(tmpfile, index=False) sample_notes_path = tmpfile.name print(f"āœ“ Sample patient notes loaded and saved to temp file: {sample_notes_path}") except Exception as e: gr.Warning(f"Could not pre-load sample patient notes: {e}") print(f"āœ— Warning: Could not pre-load sample patient notes: {e}") # Custom theme and CSS for a cleaner, modern look theme = gr.themes.Soft( primary_hue="blue", secondary_hue="slate", ).set( body_background_fill="*neutral_50", block_background_fill="white", block_border_width="1px", block_label_background_fill="*primary_50", ) custom_css = """ .gradio-container { font-family: 'Inter', Arial, sans-serif !important; } .model-status { min-height: 80px !important; font-size: 0.9em; } .status-box { background: #f9fafb; border: 1px solid #e5e7eb; border-radius: 8px; padding: 10px; } h1 { color: #1e3a8a; } """ with gr.Blocks(title="MatchMiner-AI", theme=theme, css=custom_css) as demo: with gr.Row(variant="panel"): with gr.Column(scale=4): gr.Markdown(""" # šŸ„ MatchMiner-AI **Clinical Trial Matching Pipeline** """) with gr.Column(scale=1): pass with gr.Tabs(): # ============= TAB 1: PATIENT INPUT ============= with gr.Tab("1ļøāƒ£ Patient Input"): with gr.Tab("Option A: Upload Clinical Notes"): with gr.Group(): gr.Markdown("### šŸ“„ Upload Records") notes_file = gr.File( label="Upload Patient Notes (CSV or Excel)", file_types=[".csv", ".xlsx", ".xls"], value=sample_notes_path ) # Hidden advanced option with gr.Accordion("Advanced Options", open=False): prob_threshold = gr.Slider( minimum=0.0, maximum=1.0, value=0.1, step=0.05, label="Tagger Threshold", info="Probability threshold for including excerpts" ) process_notes_btn = gr.Button("Process Notes", variant="primary") with gr.Row(): with gr.Column(): notes_status = gr.Textbox(label="Processing Status", interactive=False) with gr.Column(): pass # Removed summary button from here long_note_output = gr.Textbox( label="Extracted Patient History (Long Note)", lines=10, interactive=False ) process_notes_btn.click( fn=process_patient_notes, inputs=[notes_file, prob_threshold], outputs=[notes_status, long_note_output] ) with gr.Tab("Option B: Enter Patient Summary"): gr.Markdown("Enter a patient summary directly (skip note processing)") # Shared summary fields in a visual group with gr.Group(): gr.Markdown("### šŸ“ Patient Summary & Boilerplate") # FIX: Moved Summarize button here, right above the summary box summarize_btn = gr.Button("Summarize Patient History (from Long Note)", variant="secondary") with gr.Row(): patient_summary = gr.Textbox( label="Patient Summary", lines=12, placeholder="Enter or generate patient summary here...", info="Age, sex, Cancer type, histology, extent, biomarkers, treatment history" ) patient_boilerplate = gr.Textbox( label="Patient Boilerplate Text", lines=12, placeholder="Mentions of exclusion criteria (brain mets, etc.)", info="Evidence of potential boilerplate exclusions" ) # Wire up summarization to output to BOTH textboxes summarize_btn.click( fn=summarize_patient_history, inputs=[long_note_output], outputs=[patient_summary, patient_boilerplate] ) # ============= TAB 2: TRIAL DATABASE ============= with gr.Tab("2ļøāƒ£ Trial Database"): with gr.Row(): with gr.Column(scale=1): gr.Markdown(""" ### šŸ—ƒļø Upload Trial Space Database Upload a CSV or Excel file containing trial information. **Required Columns:** `nct_id`, `this_space`, `trial_text`, `trial_boilerplate_text` **šŸ’” TIP:** For faster loading, use pre-embedded trials defined in `config.py`. """) trial_file = gr.File( label="Upload Trial Database", file_types=[".csv", ".xlsx", ".xls"] ) trial_upload_btn = gr.Button("Load and Embed Trials", variant="primary") trial_status = gr.Textbox( label="Status", interactive=False, value=state.auto_load_status.get("trials", "") ) with gr.Column(scale=2): gr.Markdown("### Preview") trial_preview = gr.Dataframe( label="Preview (first 10 trials)", interactive=False, value=state.trial_preview_df, wrap=True ) trial_upload_btn.click( fn=load_and_embed_trials, inputs=[trial_file], outputs=[trial_status, trial_preview] ) # ============= TAB 3: TRIAL MATCHING ============= with gr.Tab("3ļøāƒ£ Trial Matching"): with gr.Row(): with gr.Column(scale=1): match_btn = gr.Button("šŸ” Find Matching Trials", variant="primary", size="lg") with gr.Column(scale=3): # Hidden advanced option for sliders with gr.Accordion("Search Settings", open=False): top_k_slider = gr.Slider( minimum=5, maximum=50, value=20, step=5, label="Number of Top Trials to Check", info="How many top-ranked trials to run eligibility checks on" ) gr.Markdown("### šŸ“Š Results") with gr.Row(): with gr.Column(scale=7): results_df = gr.Dataframe( label="Matched Trials", interactive=False, wrap=True, datatype=["str", "markdown", "markdown", "str", "str"], # Markdown enables colored text/emojis column_widths=["15%", "15%", "15%", "10%", "45%"] ) with gr.Column(scale=5): trial_details = gr.Markdown( label="Trial Details", value="
šŸ‘ˆ Click on a trial row to see full details here
" ) # Wire up matching match_btn.click( fn=match_trials, inputs=[patient_summary, patient_boilerplate, top_k_slider], outputs=[results_df] ) results_df.select( fn=get_trial_details, inputs=[results_df], outputs=[trial_details] ) # ============= TAB 4: MODEL CONFIGURATION ============= with gr.Tab("4ļøāƒ£ Model Configuration"): gr.Markdown("### šŸ§  Model Management") status_msg = """ **Config file detected** - Models will auto-load on startup. """ if HAS_CONFIG else """ **No config file found** - Please load models manually below. """ gr.Info(status_msg) with gr.Group(): with gr.Row(): with gr.Column(): tagger_input = gr.Textbox(label="TinyBERT Tagger Model", placeholder="prajjwal1/bert-tiny") tagger_btn = gr.Button("Load Tagger") tagger_status = gr.Textbox(label="Status", interactive=False, value=state.auto_load_status.get("tagger", ""), elem_classes=["model-status"]) with gr.Column(): embedder_input = gr.Textbox(label="Trial Space Embedder", placeholder="Qwen/Qwen3-Embedding-0.6B") embedder_btn = gr.Button("Load Embedder") embedder_status = gr.Textbox(label="Status", interactive=False, value=state.auto_load_status.get("embedder", ""), elem_classes=["model-status"]) embedder_warning = gr.Textbox(visible=False) with gr.Group(): with gr.Row(): with gr.Column(): llm_input = gr.Textbox(label="LLM Model (Summarization)", placeholder="openai/gpt-oss-120b") llm_btn = gr.Button("Load LLM") llm_status = gr.Textbox(label="Status", interactive=False, value=state.auto_load_status.get("llm", ""), elem_classes=["model-status"]) with gr.Column(): trial_checker_input = gr.Textbox(label="Trial Checker Model", placeholder="answerdotai/ModernBERT-large") trial_checker_btn = gr.Button("Load Trial Checker") trial_checker_status = gr.Textbox(label="Status", interactive=False, value=state.auto_load_status.get("trial_checker", ""), elem_classes=["model-status"]) with gr.Row(): with gr.Column(scale=1): boilerplate_checker_input = gr.Textbox(label="Boilerplate Checker Model", placeholder="answerdotai/ModernBERT-large") boilerplate_checker_btn = gr.Button("Load Boilerplate Checker") boilerplate_checker_status = gr.Textbox(label="Status", interactive=False, value=state.auto_load_status.get("boilerplate_checker", ""), elem_classes=["model-status"]) with gr.Column(scale=1): pass # Wire up model loading tagger_btn.click(fn=load_tagger_model, inputs=[tagger_input], outputs=[tagger_status, gr.Textbox(visible=False)]) embedder_btn.click(fn=load_embedder_model, inputs=[embedder_input], outputs=[embedder_status, gr.Textbox(visible=False), embedder_warning]) llm_btn.click(fn=load_llm_model, inputs=[llm_input], outputs=[llm_status, gr.Textbox(visible=False)]) trial_checker_btn.click(fn=load_trial_checker, inputs=[trial_checker_input], outputs=[trial_checker_status, gr.Textbox(visible=False)]) boilerplate_checker_btn.click(fn=load_boilerplate_checker, inputs=[boilerplate_checker_input], outputs=[boilerplate_checker_status, gr.Textbox(visible=False)]) # ============= TAB 5: TRIAL SPACE EXTRACTION ============= with gr.Tab("5ļøāƒ£ Trial Space Extraction"): gr.Markdown(""" ### šŸ§¬ Extract Trial Spaces Paste clinical trial text (title + summary + eligibility) to extract structured spaces. """) with gr.Row(): with gr.Column(): trial_text_input = gr.Textbox( label="Clinical Trial Text", placeholder="Paste text from ClinicalTrials.gov...", lines=15, ) extract_btn = gr.Button("Extract Trial Spaces", variant="primary") with gr.Column(): trial_spaces_output = gr.Textbox( label="Extracted Results", lines=15, interactive=False ) extract_btn.click( fn=extract_trial_spaces, inputs=[trial_text_input], outputs=[trial_spaces_output] ) return demo # ============================================================================ # MAIN # ============================================================================ if __name__ == "__main__": print(f"Device: {state.device}") print(f"GPU Available: {torch.cuda.is_available()}") if torch.cuda.is_available(): print(f"GPU Count: {torch.cuda.device_count()}") # Auto-load models from config if available if HAS_CONFIG: auto_load_models_from_config() # Auto-load trials after embedder is ready if state.embedder_model is not None or (hasattr(config, 'PREEMBEDDED_TRIALS') and config.PREEMBEDDED_TRIALS): auto_load_trials_from_config() demo = create_interface() demo.launch( server_name="0.0.0.0", server_port=7860, share=False )