app.py

#!/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(". ", "<excerpt break>")
    parts = [p.strip() for p in t2.split("<excerpt break>") 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: <cancer_type_allowed>. Histology allowed: <histology_allowed>. Cancer burden allowed: <cancer_burden_allowed>. "
    "Prior treatment required: <prior_treatments_requred>. Prior treatment excluded: <prior_treatments_excluded>. Biomarkers required: "
    "<biomarkers_required>. Biomarkers excluded: <biomarkers_excluded>.\n"
    "2. Cancer type allowed: <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
    )