#!/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 # <-- ADDED FOR PRE-EMBEDDING SUPPORT 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 # 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 # FIX #1: Store the trial preview for display at startup 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 ) # ============================================================================ # TRIAL SPACE EXTRACTION CONSTANTS # ============================================================================ 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 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), and 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; that " "renders a patient eligible for the trial.\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:\n" "1. 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" "After you output the trial spaces, output a newline, then the text \"Boilerplate exclusions:\", 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" ) 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." ) REASONING_MARKER = "assistantfinal" BOILERPLATE_MARKER = "Boilerplate exclusions:" # ============================================================================ # 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: if not os.path.exists(f"{config.PREEMBEDDED_TRIALS}_data.pkl"): print(f"✗ Pre-embedded trial files not found: {config.PREEMBEDDED_TRIALS}_*") state.auto_load_status["trials"] = f"✗ Pre-embedded files not found: {config.PREEMBEDDED_TRIALS}_*" else: print("\n" + "="*70) print(f"AUTO-LOADING PRE-EMBEDDED TRIALS: {config.PREEMBEDDED_TRIALS}") print("="*70) status, preview = load_preembedded_trials(config.PREEMBEDDED_TRIALS) state.auto_load_status["trials"] = status # FIX #1: 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 # FIX #1: 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=512 ) 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 = 1500 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(preembedded_prefix: str) -> Tuple[str, pd.DataFrame]: """Load pre-embedded trial database from disk.""" try: import pickle data_file = f"{preembedded_prefix}_data.pkl" vectors_file = f"{preembedded_prefix}_vectors.npy" metadata_file = f"{preembedded_prefix}_metadata.json" # Check files exist if not os.path.exists(data_file): return f"✗ Pre-embedded data file not found: {data_file}", None if not os.path.exists(vectors_file): return f"✗ Pre-embedded vectors file not found: {vectors_file}", None print(f"\n{'='*70}") print(f"LOADING PRE-EMBEDDED TRIALS") print(f"{'='*70}") print(f"Loading from: {preembedded_prefix}_*") # Load metadata if available if os.path.exists(metadata_file): with open(metadata_file, 'r') as f: metadata = json.load(f) print(f"Metadata:") print(f" Created: {metadata.get('created_at', 'unknown')}") print(f" Embedder: {metadata.get('embedder_model', 'unknown')}") print(f" Trials: {metadata.get('num_trials', 'unknown')}") print(f" Embedding dim: {metadata.get('embedding_dim', 'unknown')}") # Load dataframe print(f"Loading trial dataframe...") with open(data_file, 'rb') as f: df = pickle.load(f) print(f"✓ Loaded {len(df)} trials") # Load embeddings print(f"Loading embeddings...") embeddings = np.load(vectors_file) print(f"✓ Loaded embeddings: {embeddings.shape}") # Validate if len(df) != embeddings.shape[0]: return ( f"✗ Mismatch: {len(df)} trials but {embeddings.shape[0]} embeddings", None ) # Store in state state.trial_spaces_df = df state.trial_embeddings = embeddings print(f"{'='*70}") print(f"PRE-EMBEDDED TRIALS LOADED SUCCESSFULLY") print(f"{'='*70}\n") preview = df[['nct_id', 'this_space']].head(10) return f"✓ Loaded {len(df)} pre-embedded trials from {preembedded_prefix}_*", 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=1500) ) # 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.5) -> 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() 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['date_text'] = ( keep['date_str'] + " " + keep['note_type'] + " " + keep['excerpt'] ) # Create long note long_note = "\n".join(keep['date_text'].tolist()) stats = ( f"Processed {len(df)} notes → {len(all_excerpts)} excerpts → " f"{len(keep)} relevant excerpts (threshold={prob_threshold})" ) return stats, long_note except Exception as e: return f"✗ Error processing notes: {str(e)}", "" # FIX #2: Modified function to return both summary and boilerplate as separate outputs 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 = 115000 # 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': f"""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. Document the 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. Format your response as free text, not as a table. 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: 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: {patient_text} Now, 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) # Extract just the assistant response if "assistant" in output: output = output.split("assistant")[-1] # Clean up reasoning markers if present if "assistantfinal" in output: output = output.split("assistantfinal", 1)[-1] output = output.strip() # FIX #2: Split the output into summary and boilerplate sections if "Boilerplate:" in output: parts = output.split("Boilerplate:", 1) summary_text = parts[0].strip() boilerplate_text = parts[1].strip() else: # If no boilerplate section found, put everything in summary summary_text = output boilerplate_text = "No boilerplate exclusion criteria documented" 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 "assistant" in output: output = output.split("assistant")[-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=1500) 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=2048, 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=2048, 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) # Select columns for display display_cols = [ 'nct_id', 'eligibility_probability', 'exclusion_probability', 'similarity_score', 'this_space' ] result_df = top_trials[display_cols].reset_index(drop=True) # Convert probability columns to strings with fixed decimal places # This ensures Gradio displays them correctly without extra decimals result_df['eligibility_probability'] = result_df['eligibility_probability'].apply(lambda x: f"{x:.2f}") result_df['exclusion_probability'] = result_df['exclusion_probability'].apply(lambda x: f"{x:.2f}") result_df['similarity_score'] = result_df['similarity_score'].apply(lambda x: f"{x:.3f}") 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] nct_id = df.iloc[row_idx]['nct_id'] this_space = df.iloc[row_idx]['this_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(): # Custom CSS for Arial font and styling custom_css = """ * { font-family: Arial, sans-serif !important; } .model-status { min-height: 120px !important; } """ with gr.Blocks(title="MatchMiner-AI Demo App", theme=gr.themes.Soft(), css=custom_css) as demo: gr.Markdown(""" # 🏥 Clinical Trial Matching Pipeline Match cancer patients to relevant clinical trials using open-source AI-powered eligibility pre-screening. """) with gr.Tabs(): # ============= TAB 1: PATIENT INPUT ============= with gr.Tab("1️⃣ Patient Input"): gr.Markdown("### Patient Data Entry") with gr.Tab("Option A: Upload Clinical Notes"): gr.Markdown(""" Upload patient clinical notes as CSV or Excel with columns: - `date`: Date of note - `text`: Note text - `note_type` (optional): Type of note """) notes_file = gr.File( label="Upload Patient Notes (CSV or Excel)", file_types=[".csv", ".xlsx", ".xls"] ) prob_threshold = gr.Slider( minimum=0.0, maximum=1.0, value=0.5, step=0.05, label="Tagger Threshold", info="Probability threshold for including excerpts" ) process_notes_btn = gr.Button("Process Notes", variant="primary", size="lg") notes_status = gr.Textbox(label="Processing Status", interactive=False) long_note_output = gr.Textbox( label="Extracted Patient History (Long Note)", lines=10, interactive=False ) summarize_btn = gr.Button("Summarize Patient History", variant="secondary", size="lg") 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 gr.Markdown("### Patient Summary") patient_summary = gr.Textbox( label="Patient Summary", lines=15, placeholder="Enter or generate patient summary here...", info="Cancer type, histology, extent, biomarkers, treatment history" ) patient_boilerplate = gr.Textbox( label="Patient Boilerplate Text", lines=5, placeholder="Mentions of exclusion criteria (brain mets, etc.)", info="Evidence of potential boilerplate exclusions" ) # FIX #2: 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"): gr.Markdown(""" ### Upload Trial Space Database Upload a CSV or Excel file containing trial information with these required columns: - `nct_id`: NCT identifier - `this_space`: Summary of eligibility criteria - `trial_text`: Full trial description - `trial_boilerplate_text`: Boilerplate exclusion criteria **💡 TIP:** For faster loading, use pre-embedded trials! See instructions in config_example.py **⚠️ Note:** If you change the embedder model, trials will need to be re-embedded. """) trial_file = gr.File( label="Upload Trial Database (CSV or Excel)", file_types=[".csv", ".xlsx", ".xls"] ) trial_upload_btn = gr.Button("Load and Embed Trials", variant="primary", size="lg") trial_status = gr.Textbox( label="Status", interactive=False, value=state.auto_load_status.get("trials", "") ) # FIX #1: Set the initial value to the preview from auto-loading trial_preview = gr.Dataframe( label="Preview (first 10 trials)", interactive=False, value=state.trial_preview_df, column_widths=["20%", "80%"] # NCT ID wider, this_space takes remaining space ) 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"): gr.Markdown("### Match Patient to Trials") 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" ) match_btn = gr.Button("🔍 Find Matching Trials", variant="primary", size="lg") gr.Markdown(""" ### Results Click on a row in the table to see full trial details on the right. **Columns:** - **Eligibility Probability**: Predicted likelihood this trial is reasonable for the patient (higher is better) - **Exclusion Probability**: Predicted likelihood the patient fails boilerplate exclusions (lower is better) - **Similarity Score**: Embedding similarity between patient and trial """) with gr.Row(): with gr.Column(scale=1): results_df = gr.Dataframe( label="Matched Trials", interactive=False, wrap=True, column_widths=["20%", "15%", "15%", "12%", "38%"] # Adjusted for side-by-side layout ) with gr.Column(scale=1): trial_details = gr.Markdown( label="Trial Details", value="👈 Click on a trial in the table to see its 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("### Load Required Models") if HAS_CONFIG: gr.Markdown(""" ✓ **Config file detected** - Models will auto-load on startup. You can still manually load different models below if needed. """) else: gr.Markdown(""" ℹ️ **No config file found** - Load models manually below. To enable auto-loading, create a `config.py` file (see `config_example.py`). """) with gr.Row(): with gr.Column(): tagger_input = gr.Textbox( label="TinyBERT Tagger Model", placeholder="prajjwal1/bert-tiny or path/to/model", info="Model for extracting relevant excerpts from clinical notes" ) tagger_btn = gr.Button("Load Tagger", variant="primary") tagger_status = gr.Textbox( label="Status", interactive=False, elem_classes=["model-status"], value=state.auto_load_status.get("tagger", "") ) with gr.Column(): embedder_input = gr.Textbox( label="Trial Space Embedder Model", placeholder="Qwen/Qwen3-Embedding-0.6B or path/to/reranker_round2.model", info="Sentence transformer for embedding patient summaries and trials" ) embedder_btn = gr.Button("Load Embedder", variant="primary") embedder_status = gr.Textbox( label="Status", interactive=False, elem_classes=["model-status"], value=state.auto_load_status.get("embedder", "") ) embedder_warning = gr.Textbox(label="", interactive=False, visible=False) with gr.Row(): with gr.Column(): llm_input = gr.Textbox( label="LLM Model (for Summarization)", placeholder="openai/gpt-oss-120b or path/to/model", info="Large language model for summarizing patient histories" ) llm_btn = gr.Button("Load LLM", variant="primary") llm_status = gr.Textbox( label="Status", interactive=False, elem_classes=["model-status"], value=state.auto_load_status.get("llm", "") ) with gr.Column(): trial_checker_input = gr.Textbox( label="Trial Checker Model", placeholder="answerdotai/ModernBERT-large or path/to/modernbert-trial-checker", info="ModernBERT model for eligibility prediction" ) trial_checker_btn = gr.Button("Load Trial Checker", variant="primary") trial_checker_status = gr.Textbox( label="Status", interactive=False, elem_classes=["model-status"], value=state.auto_load_status.get("trial_checker", "") ) with gr.Row(): with gr.Column(): boilerplate_checker_input = gr.Textbox( label="Boilerplate Checker Model", placeholder="answerdotai/ModernBERT-large or path/to/modernbert-boilerplate-checker", info="ModernBERT model for boilerplate exclusion prediction" ) boilerplate_checker_btn = gr.Button("Load Boilerplate Checker", variant="primary") boilerplate_checker_status = gr.Textbox( label="Status", interactive=False, elem_classes=["model-status"], value=state.auto_load_status.get("boilerplate_checker", "") ) # 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 from Clinical Trial Text This tool extracts structured trial spaces and boilerplate exclusion criteria from clinical trial documents. **Instructions:** 1. Copy and paste the clinical trial text (title + summary + eligibility criteria) from ClinicalTrials.gov into the text box below 2. Click "Extract Trial Spaces" to process the text with the LLM 3. Review the extracted spaces and boilerplate criteria in the output **Note:** The LLM model must be loaded first (see Model Configuration tab). """) with gr.Row(): with gr.Column(): trial_text_input = gr.Textbox( label="Clinical Trial Text", placeholder="Paste the concatenation of clinical trial title, summary, and eligibility criteria here...", lines=15, max_lines=20 ) extract_btn = gr.Button("Extract Trial Spaces", variant="primary", size="lg") with gr.Column(): trial_spaces_output = gr.Textbox( label="Extracted Trial Spaces and Boilerplate Criteria", lines=15, max_lines=20, interactive=False ) gr.Markdown(""" ### About Trial Spaces A **trial space** is a unique combination of: - Cancer primary site and histology - Required and excluded prior treatments - Cancer burden (e.g., metastatic disease) - Required and excluded tumor biomarkers **Boilerplate exclusions** are common trial exclusion criteria such as: - History of pneumonitis, heart failure, renal/liver dysfunction - Uncontrolled brain metastases - HIV or hepatitis infection - Poor performance status """) # Wire up extraction extract_btn.click( fn=extract_trial_spaces, inputs=[trial_text_input], outputs=[trial_spaces_output] ) gr.Markdown(""" --- ### Instructions 1. **Model Configuration**: Load required models (auto-loads from config.py if present) 2. **Trial Database**: Upload trial CSV/Excel OR use pre-embedded trials (much faster!) 3. **Patient Input**: Upload clinical notes OR enter a patient summary directly 4. **Trial Matching**: Click to find and rank matching trials with eligibility predictions **💡 PRO TIP:** Use `preembed_trials.py` to pre-embed your trial database for 10-100x faster loading! **Note**: First-time model loading may take a few minutes. GPU acceleration is used if available. """) 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 )