Add all BioGPT app files with LFS tracking

Dockerfile

@@ -0,0 +1,36 @@
+# Use an official Python runtime as a parent image (e.g., 3.11 or 3.12)
+FROM python:3.11-slim
+
+# Set the working directory in the container
+WORKDIR /app
+
+# Install system dependencies if needed (e.g., build-essential if compiling fastBPE)
+# RUN apt-get update && apt-get install -y --no-install-recommends build-essential && rm -rf /var/lib/apt/lists/*
+
+# Copy the requirements file into the container first (for Docker caching)
+COPY requirements.txt requirements.txt
+
+# Install Python dependencies
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy the application code and necessary files into the container
+# IMPORTANT: This assumes 'fast' is pre-compiled for Linux x86_64
+COPY app.py .
+COPY hoc_best.pt .
+COPY bpecodes .
+COPY dict.txt .
+COPY fast .       # Copy the 'fast' executable directly
+COPY templates/ ./templates/
+
+# Ensure the 'fast' binary is executable within the container
+RUN chmod +x ./fast
+
+# Make port 5000 available (as defined in README.md and app.py)
+EXPOSE 5000
+
+# Define environment variable for Python output buffering
+ENV PYTHONUNBUFFERED=1
+
+# Run the app using waitress (a production-ready WSGI server)
+# Assumes your Flask app instance is named 'app' in 'app.py'
+CMD ["waitress-serve", "--host=0.0.0.0", "--port=5000", "app:app"]

app.py

@@ -0,0 +1,1067 @@
+import os
+import subprocess
+import math
+import difflib
+import tempfile
+import numpy as np
+import pandas as pd
+from tqdm.auto import tqdm
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import dataclass
+from sacremoses import MosesDetokenizer
+from flask import Flask, request, jsonify, render_template
+import traceback # Import traceback for better error logging
+import re # Import the regular expression module
+
+# --- Constants and Paths ---
+# Ensure these files are in the same directory as app.py or provide correct paths
+FINETUNED_MODEL_PATH = "hoc_best.pt"
+BPE_CODES_PATH = "bpecodes"
+DICT_TXT_PATH = "dict.txt"
+FASTBPE_BIN_PATH = "./fast" # Assumes fast executable is alongside app.py
+
+HALLMARKS = [ # Keep this consistent with training/evaluation
+    "activating invasion and metastasis", "avoiding immune destruction",
+    "cellular energetics", "enabling replicative immortality",
+    "evading growth suppressors", "genomic instability and mutation",
+    "inducing angiogenesis", "resisting cell death",
+    "sustaining proliferative signaling", "tumor promoting inflammation",
+]
+
+# --- Model Architecture Definitions (Copy from your notebook) ---
+# NOTE: Make sure these classes are IDENTICAL to the ones used for training
+#       including GPTConfig, LayerNorm, CausalSelfAttention, MLP, Block, GPT, GPTWithSoftPrompt
+
+@dataclass
+class GPTConfig:
+    block_size: int
+    vocab_size: int
+    n_layer: int
+    n_head: int
+    n_embd: int
+    dropout: float = 0.0
+    bias: bool = True
+
+class LayerNorm(nn.Module):
+    # (Copied from notebook)
+    def __init__(self, ndim, bias):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(ndim))
+        self.bias = nn.Parameter(torch.zeros(ndim)) if bias else None
+    def forward(self, x):
+        return F.layer_norm(x, self.weight.shape, self.weight, self.bias, 1e-5)
+
+class CausalSelfAttention(nn.Module):
+    # (Copied from notebook)
+    def __init__(self, config):
+        super().__init__()
+        assert config.n_embd % config.n_head == 0
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd, bias=config.bias)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd, bias=config.bias)
+        self.attn_dropout = nn.Dropout(config.dropout)
+        self.resid_dropout = nn.Dropout(config.dropout)
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.flash = hasattr(F, 'scaled_dot_product_attention') # Check for flash attention
+        if not self.flash:
+            # print("Warning: Flash Attention not available.") # Optional warning
+            # Make the buffer persistent otherwise device mismatches during forward pass
+            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
+                                        .view(1, 1, config.block_size, config.block_size), persistent=True)
+        #else:
+            # print("Using Flash Attention.") # Optional info
+
+    def forward(self, x):
+        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)
+        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
+        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
+        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
+
+        # causal self-attention; Self-attend: (B, nh, T, hs) x (B, nh, hs, T) -> (B, nh, T, T)
+        if self.flash:
+            # efficient attention using Flash Attention CUDA kernels
+            y = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=self.attn_dropout.p if self.training else 0.0, is_causal=True)
+        else:
+            # manual implementation of attention
+            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+            # Ensure bias buffer is used correctly
+            # Check if bias buffer exists before using it
+            if hasattr(self, 'bias'):
+                att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
+            else:
+                 # Fallback if somehow bias wasn't registered (shouldn't happen with persistent=True)
+                 mask = torch.tril(torch.ones(T, T, device=x.device)).view(1, 1, T, T)
+                 att = att.masked_fill(mask == 0, float('-inf'))
+
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
+
+        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side
+        # output projection
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+
+class MLP(nn.Module):
+     # (Copied from notebook)
+    def __init__(self, config):
+        super().__init__()
+        self.c_fc    = nn.Linear(config.n_embd, 4 * config.n_embd, bias=config.bias)
+        self.gelu    = nn.GELU()
+        self.c_proj  = nn.Linear(4 * config.n_embd, config.n_embd, bias=config.bias)
+        self.dropout = nn.Dropout(config.dropout)
+    def forward(self, x):
+        x = self.c_fc(x)
+        x = self.gelu(x)
+        x = self.c_proj(x)
+        x = self.dropout(x)
+        return x
+
+class Block(nn.Module):
+     # (Copied from notebook)
+    def __init__(self, config):
+        super().__init__()
+        self.ln1 = LayerNorm(config.n_embd, bias=config.bias)
+        self.attn = CausalSelfAttention(config)
+        self.ln2 = LayerNorm(config.n_embd, bias=config.bias)
+        self.mlp = MLP(config)
+    def forward(self, x):
+        x = x + self.attn(self.ln1(x))
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+class GPT(nn.Module):
+     # (Copied from notebook - simplified _init_weights and removed generate)
+    def __init__(self, config):
+        super().__init__()
+        #assert config.vocab_size is not None
+        #assert config.block_size is not None
+        self.config = config
+
+        self.transformer = nn.ModuleDict(dict(
+            wte = nn.Embedding(config.vocab_size, config.n_embd),
+            wpe = nn.Embedding(config.block_size, config.n_embd),
+            drop = nn.Dropout(config.dropout),
+            h = nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
+            ln_f = LayerNorm(config.n_embd, bias=config.bias),
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.transformer.wte.weight = self.lm_head.weight # https://paperswithcode.com/method/weight-tying
+
+        # init all weights
+        self.apply(self._init_weights)
+        # apply special scaled init to the residual projections, per GPT-2 paper
+        for pn, p in self.named_parameters():
+            if pn.endswith('c_proj.weight'):
+                torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                torch.nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            torch.nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None):
+        device = idx.device
+        b, t = idx.size()
+        if t > self.config.block_size:
+             # Crop sequence if longer than block size
+             print(f"Warning: Input sequence length ({t}) > block size ({self.config.block_size}). Cropping.")
+             idx = idx[:, -self.config.block_size:]
+             t = self.config.block_size
+        #assert t <= self.config.block_size, f"Cannot forward sequence of length {t}, block size is only {self.config.block_size}"
+        pos = torch.arange(0, t, dtype=torch.long, device=device) # shape (t)
+
+        # forward the GPT model itself
+        tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
+        pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
+        x = self.transformer.drop(tok_emb + pos_emb)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+
+        if targets is not None:
+            # if we are given some desired targets also calculate the loss
+            logits = self.lm_head(x)
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1)
+            return logits, loss
+        else:
+            # inference-time mini-optimization: only forward the lm_head on the very last position
+            logits = self.lm_head(x[:, [-1], :]) # note: using list [-1] to preserve the time dim
+            # Check for NaN/Inf in logits before returning
+            if torch.isnan(logits).any() or torch.isinf(logits).any():
+                print("WARNING: NaN or Inf detected in logits during inference.")
+                # Handle appropriately - maybe return an error indicator or zero logits?
+                # For now, just print warning.
+            return logits, None
+
+class GPTWithSoftPrompt(nn.Module):
+     # (Copied from notebook - simplified)
+    def __init__(self, base_gpt: GPT, prompt_len=1):
+        super().__init__()
+        self.config = base_gpt.config
+        self.transformer = base_gpt.transformer
+        self.lm_head = base_gpt.lm_head
+        C = self.config.n_embd
+        self.soft_prompt = nn.Parameter(torch.zeros(1, prompt_len, C)) # Keep on CPU first
+        nn.init.normal_(self.soft_prompt, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None):
+        B, T = idx.shape
+        device = idx.device # Get device from input tensor
+
+        # Make sure soft_prompt is on the same device as input
+        soft_prompt_on_device = self.soft_prompt.to(device)
+
+        # token + pos
+        tok_emb = self.transformer.wte(idx) # (B,T,C)
+        pos = torch.arange(0, T, dtype=torch.long, device=device)
+        pos_emb = self.transformer.wpe(pos) # (T,C)
+        x_tokens = tok_emb + pos_emb
+
+        # prepend soft prompt
+        soft = soft_prompt_on_device.expand(B, -1, -1) # (B,P,C)
+
+        # --- FIX: Define P before the if/else block ---
+        P = soft.size(1) # Get soft prompt length
+
+        x = torch.cat([soft, x_tokens], dim=1) # (B,P+T,C)
+
+        # --- Standard Transformer forward pass ---
+        x = self.transformer.drop(x)
+        for block in self.transformer.h:
+            x = block(x)
+        x = self.transformer.ln_f(x)
+        logits = self.lm_head(x) # (B,P+T,V)
+        # --- End Standard ---
+
+        if targets is None:
+            # Inference: return logits for the last token of the *original* sequence
+            # We need the prediction *after* the last input token, which is at index T (P+T-1 overall)
+            # Use P which is now defined
+            # Ensure index is within bounds
+            target_logit_index = P + T - 1
+            if target_logit_index >= logits.size(1):
+                 print(f"Warning: Calculated logit index {target_logit_index} out of bounds for logits shape {logits.shape}. Returning last logit.")
+                 target_logit_index = -1 # Fallback to last logit
+
+            final_logits = logits[:, target_logit_index, :]
+            # Check for NaN/Inf
+            if torch.isnan(final_logits).any() or torch.isinf(final_logits).any():
+                 print(f"WARNING: NaN or Inf detected in final_logits at index {target_logit_index}.")
+                 # Handle appropriately - maybe return zeros or raise an error?
+                 # For now, just print warning. Let the calling function handle it.
+
+            return final_logits, None # Return (B, V)
+        else:
+            # Training loss calculation (copied from notebook)
+            # P is already defined above
+            pad_ignore = torch.full((B, P), -1, dtype=targets.dtype, device=device)
+            full_targets = torch.cat([pad_ignore, targets], dim=1)
+            logits_lm  = logits[:, :-1, :].contiguous()
+            targets_lm = full_targets[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                logits_lm.view(-1, logits_lm.size(-1)),
+                targets_lm.view(-1),
+                ignore_index=-1
+            )
+            # Check for NaN/Inf in loss
+            if torch.isnan(loss) or torch.isinf(loss):
+                print("WARNING: NaN or Inf detected in loss calculation.")
+                # Potentially add debugging info here (e.g., print shapes, inputs)
+            return logits, loss
+
+    # --- Constrained generation method (from Section 2.9) ---
+    @torch.no_grad()
+    def generate_labels(self, idx, allowed_mask, max_new_tokens=24, temperature=0.0):
+        self.eval() # Ensure model is in eval mode
+        B = idx.size(0)
+        # Add soft prompt length to effective block size consideration
+        P = self.soft_prompt.size(1)
+        # Correct effective block size based on GPT class logic
+        effective_block_size = self.config.block_size # GPT forward handles cropping
+
+        # Start with input index
+        out = idx.clone() # Clone to avoid modifying original input
+
+        # Ensure allowed_mask is on the correct device
+        allowed_mask = allowed_mask.to(idx.device)
+
+        finished = torch.zeros(B, dtype=torch.bool, device=idx.device)
+
+        # Get global eos_id safely
+        global eos_id
+        current_eos_id = eos_id # Use the globally loaded eos_id
+
+        for step in range(max_new_tokens):
+            # Crop context if it exceeds model's block size (GPT forward handles this internally now)
+            # ctx = out if out.size(1) <= effective_block_size else out[:, -effective_block_size:]
+            ctx = out # Pass the current sequence
+
+            # Forward pass - expects shape (B, T), model handles soft prompt internally
+            # It returns logits for the *next* token prediction after the last token in ctx
+            logits, _ = self(ctx) # Gets logits for last token prediction, shape (B, V)
+
+            # Check for NaN/Inf in logits
+            if torch.isnan(logits).any() or torch.isinf(logits).any():
+                print(f"WARNING: NaN or Inf detected in logits during generation step {step}. Stopping generation.")
+                # Return what we have so far, excluding potentially bad last token
+                return out[:, idx.size(1):] # Or handle error differently
+
+            # Apply constraint mask
+            # Ensure mask shape matches logits shape
+            if logits.shape != allowed_mask.shape:
+                 print(f"Warning: Logits shape {logits.shape} doesn't match mask shape {allowed_mask.shape}. Reshaping mask.")
+                 # This assumes mask needs batch dim added
+                 current_mask = allowed_mask.unsqueeze(0).expand_as(logits)
+            else:
+                 current_mask = allowed_mask
+
+            logits = logits + current_mask
+
+            # Sample next token
+            if temperature <= 0:
+                # Greedy decoding
+                next_id = torch.argmax(logits, dim=-1)  # (B,)
+            else:
+                # Temperature sampling
+                probs = F.softmax(logits / temperature, dim=-1)
+                # Check for NaN/Inf in probs
+                if torch.isnan(probs).any() or torch.isinf(probs).any():
+                     print(f"WARNING: NaN or Inf detected in probabilities during generation step {step}. Using argmax fallback.")
+                     next_id = torch.argmax(logits, dim=-1) # Fallback to greedy
+                else:
+                    try:
+                         next_id = torch.multinomial(probs, num_samples=1).squeeze(1) # (B,)
+                    except RuntimeError as e:
+                         print(f"WARNING: torch.multinomial failed: {e}. Using argmax fallback.")
+                         next_id = torch.argmax(logits, dim=-1) # Fallback to greedy
+
+
+            # Handle finished sequences (force EOS) and update output
+            # Check if current_eos_id is valid
+            if not isinstance(current_eos_id, int):
+                print(f"Warning: Global eos_id is not an integer ({current_eos_id}). Defaulting to 0.")
+                current_eos_id = 0
+            next_id = next_id.masked_fill(finished, current_eos_id) # Use the validated eos_id
+
+            # Check if next_id contains invalid values (e.g., negative)
+            if (next_id < 0).any():
+                 print(f"WARNING: Negative token ID generated: {next_id}. Clipping to 0.")
+                 next_id = torch.clamp(next_id, min=0)
+
+
+            # Append the next token ID
+            out = torch.cat([out, next_id.unsqueeze(1)], dim=1)
+
+            # Update finished status
+            finished |= (next_id == current_eos_id)
+
+            # Stop if all sequences in the batch are finished
+            if bool(finished.all()):
+                # print(f"Generation finished early at step {step+1}") # Optional debug info
+                break
+        # else:
+            # print(f"Generation reached max_new_tokens ({max_new_tokens})") # Optional debug info
+
+        # Return only the generated part (excluding the initial idx length)
+        return out[:, idx.size(1):]
+
+
+# --- Tokenizer Helper Functions ---
+# Added robustness and error checks
+# Global tokenizer maps and special IDs, loaded once at startup
+token2id, id2token = {}, {}
+eos_id = 0 # Default, will be overwritten
+pad_id = 0 # Default, will be overwritten
+detokenizer = None
+
+def load_tokenizer_data(dict_path):
+    global token2id, id2token, eos_id, pad_id, detokenizer
+    print(f"Loading vocabulary from {dict_path}...")
+    local_token2id, local_id2token = {}, {}
+    try:
+        with open(dict_path, encoding="utf-8") as f:
+            for i, line in enumerate(f):
+                parts = line.split() # Split by whitespace
+                if not parts: continue # Skip empty lines
+                tok = parts[0]
+                if tok in local_token2id:
+                     print(f"Warning: Duplicate token '{tok}' found at line {i+1}. Keeping first occurrence.")
+                     continue
+                local_token2id[tok] = i
+                local_id2token[i] = tok
+
+        # Assign to global variables only after successful loading
+        token2id = local_token2id
+        id2token = local_id2token
+
+        # Use a known special token ID if </s> is missing, otherwise default might be wrong
+        # Try multiple common EOS tokens
+        possible_eos = ["</s>", "<|endoftext|>", "[EOS]"]
+        found_eos = False
+        for eos_tok in possible_eos:
+            if eos_tok in token2id:
+                eos_id = token2id[eos_tok]
+                found_eos = True
+                print(f"Found EOS token '{eos_tok}' with ID: {eos_id}")
+                break
+        if not found_eos:
+             # If no common EOS found, fall back to the highest index or 0
+             eos_id = max(token2id.values()) if token2id else 0
+             print(f"Warning: Standard EOS tokens not found. Using highest index ({eos_id}) as EOS ID.")
+
+        # Assign pad_id, often same as eos_id or a specific <pad> token
+        pad_id = token2id.get("<pad>", eos_id) # Prefer <pad> if exists, else use eos_id
+        print(f"Using PAD ID: {pad_id}")
+
+        detokenizer = MosesDetokenizer(lang='en') # Initialize once
+        print(f"Vocabulary loaded. Size: {len(token2id)}")
+        if not detokenizer:
+             raise ValueError("MosesDetokenizer failed to initialize.")
+
+    except FileNotFoundError:
+        print(f"ERROR: Vocabulary file not found at {dict_path}")
+        raise
+    except Exception as e:
+        print(f"ERROR: Failed to load tokenizer data from {dict_path}: {e}")
+        raise
+
+def bpe_encode_lines(lines, shard_size=500, desc="BPE Encode"):
+    """ Encodes lines using external fastBPE binary. Added error checking. """
+    global BPE_CODES_PATH, FASTBPE_BIN_PATH
+    # --- Input Validation ---
+    if not isinstance(lines, list):
+         print(f"Warning: bpe_encode_lines expected a list, got {type(lines)}. Attempting conversion.")
+         try:
+              lines = list(lines)
+         except TypeError:
+              raise ValueError("Input 'lines' must be a list or convertible to a list.")
+
+    if not lines: return []
+
+    # --- Path and Executable Checks ---
+    abs_fastbpe_path = os.path.abspath(FASTBPE_BIN_PATH)
+    abs_bpe_codes_path = os.path.abspath(BPE_CODES_PATH)
+
+    if not os.path.exists(abs_fastbpe_path):
+        raise FileNotFoundError(f"fastBPE executable not found at {abs_fastbpe_path}")
+    if not os.path.exists(abs_bpe_codes_path):
+        raise FileNotFoundError(f"BPE codes file not found at {abs_bpe_codes_path}")
+    if not os.access(abs_fastbpe_path, os.X_OK):
+        print(f"Warning: fastBPE binary at {abs_fastbpe_path} is not executable. Attempting chmod...")
+        try:
+            os.chmod(abs_fastbpe_path, 0o755)
+        except OSError as e:
+            raise PermissionError(f"Failed to make fastBPE executable: {e}. Please check permissions.")
+
+    out_tokens = []
+    # Process in chunks
+    with tempfile.TemporaryDirectory() as td:
+        for start in range(0, len(lines), shard_size):
+            chunk = lines[start:start+shard_size]
+            src_path = os.path.join(td, f"src_{start}.txt")
+            dst_path = os.path.join(td, f"dst_{start}.bpe")
+
+            try:
+                # Write chunk to temp file, ensuring strings
+                with open(src_path, "w", encoding="utf-8") as f:
+                    for s in chunk:
+                        f.write(str(s or "").strip() + "\n") # Ensure string conversion
+
+                # Run fastBPE
+                cmd = [abs_fastbpe_path, "applybpe", dst_path, src_path, abs_bpe_codes_path]
+                # print(f"Running command: {' '.join(cmd)}") # Debug command
+                process = subprocess.run(
+                    cmd,
+                    capture_output=True, text=True, check=False # Don't check=True here, handle error below
+                )
+
+                # Check for errors specifically
+                if process.returncode != 0:
+                     # Log more details on failure
+                     print(f"ERROR: fastBPE failed (exit code {process.returncode}) on chunk starting at index {start}.")
+                     print(f"Command: {' '.join(cmd)}")
+                     print(f"Stderr:\n{process.stderr}")
+                     # Optionally print some input data
+                     print(f"First line of input chunk: {chunk[0] if chunk else 'N/A'}")
+                     raise subprocess.CalledProcessError(process.returncode, cmd, output=process.stdout, stderr=process.stderr)
+
+                # Read results if successful
+                with open(dst_path, "r", encoding="utf-8") as f:
+                    for line in f:
+                        out_tokens.append(line.strip().split())
+
+            except subprocess.CalledProcessError as e:
+                # Handle specific subprocess errors (already printed details)
+                raise # Re-raise to stop execution
+            except Exception as e:
+                print(f"ERROR: Unexpected error during BPE encoding chunk starting at index {start}: {e}")
+                traceback.print_exc() # Print full traceback for unexpected errors
+                raise # Re-raise
+    return out_tokens
+
+
+def tokens_to_ids(bpe_tokens):
+    """ Converts BPE token strings to IDs using the global map. Added checks. """
+    global token2id, pad_id
+    if not isinstance(bpe_tokens, list):
+         raise ValueError(f"Input 'bpe_tokens' must be a list, got {type(bpe_tokens)}.")
+
+    ids = []
+    oov_count = 0
+    for t in bpe_tokens:
+        if not isinstance(t, str):
+             print(f"Warning: Non-string token found in bpe_tokens: {t}. Using pad_id.")
+             ids.append(pad_id)
+             oov_count += 1
+             continue
+
+        id_val = token2id.get(t, pad_id)
+        ids.append(id_val)
+        if id_val == pad_id and t not in token2id:
+            oov_count += 1
+            # print(f"Warning: OOV token '{t}' mapped to pad_id {pad_id}") # Reduce noise
+    if oov_count > 0:
+         print(f"Info: Found {oov_count} OOV tokens in sequence of length {len(bpe_tokens)}.")
+    return ids, oov_count
+
+def ids_to_tokens(ids):
+    """ Converts IDs back to token strings. Added checks. """
+    global id2token
+    if not isinstance(ids, list):
+        raise ValueError(f"Input 'ids' must be a list, got {type(ids)}.")
+
+    tokens = []
+    for i in ids:
+        # Ensure ID is a valid integer before lookup
+        try:
+             # Handle potential floats or NaNs from generation
+             if isinstance(i, float) and math.isnan(i):
+                 token = "<nan>"
+             else:
+                 int_i = int(i)
+                 token = id2token.get(int_i, "<unk>")
+        except (ValueError, TypeError):
+             print(f"Warning: Could not convert ID '{i}' to int. Using '<unk>'.")
+             token = "<unk>"
+        tokens.append(token)
+    return tokens
+
+
+def bpe_decode_tokens(bpe_tokens):
+    """ Converts BPE token strings back to readable text. Added checks. """
+    global detokenizer
+    if detokenizer is None:
+        raise RuntimeError("Detokenizer not initialized. Call load_tokenizer_data first.")
+    if not isinstance(bpe_tokens, list):
+         raise ValueError(f"Input 'bpe_tokens' must be a list, got {type(bpe_tokens)}.")
+
+    # Ensure all items are strings before joining
+    try:
+        str_tokens = [str(t) for t in bpe_tokens]
+    except Exception as e:
+        print(f"Error converting tokens to strings: {e}. Tokens: {bpe_tokens}")
+        return "<decoding error>"
+
+    s = ' '.join(str_tokens).replace('@@ ', '')
+    try:
+        # Detokenizer might fail on empty or unusual input
+        return detokenizer.detokenize(s.split()) if s.strip() else ""
+    except Exception as e:
+        print(f"Error during detokenization: {e}. Input string: '{s}'")
+        return "<detokenization error>"
+
+
+# --- Prediction Helper Functions ---
+
+def to_canonical(pred_chunk: str):
+    """ Maps a predicted text chunk to a canonical hallmark name. Added checks. """
+    global HALLMARKS
+    # Ensure input is a string
+    if not isinstance(pred_chunk, str):
+         # print(f"Warning: to_canonical received non-string input: {pred_chunk}. Returning None.")
+         return None
+
+    s = pred_chunk.strip().lower()
+    low = [L.lower() for L in HALLMARKS]
+    if not s: return None
+
+    if s in low:
+        return HALLMARKS[low.index(s)]
+
+    # Use difflib for fuzzy matching
+    try:
+        best = difflib.get_close_matches(s, low, n=1, cutoff=0.7)
+        return HALLMARKS[low.index(best[0])] if best else None
+    except Exception as e:
+        print(f"Error during difflib matching for '{s}': {e}")
+        return None # Return None on error
+
+def build_allowed_token_mask(vocab_size, device):
+    """ Builds the mask for constrained decoding. Added error checks. """
+    global HALLMARKS, token2id, eos_id, pad_id
+    allowed = set()
+
+    # --- Input Validation ---
+    if vocab_size <= 0:
+        raise ValueError("Vocabulary size must be positive.")
+    if not token2id:
+        raise RuntimeError("Tokenizer vocabulary (token2id) not loaded.")
+
+    print("Encoding hallmarks for mask...")
+    try:
+        # Ensure HALLMARKS is a list of strings
+        if not isinstance(HALLMARKS, list) or not all(isinstance(h, str) for h in HALLMARKS):
+             raise ValueError("HALLMARKS must be a list of strings.")
+        hallmark_bpes = bpe_encode_lines(HALLMARKS, desc="BPE Hallmarks (for mask)")
+        for bpe_list in hallmark_bpes:
+            ids, _ = tokens_to_ids(bpe_list)
+            allowed.update(ids)
+        print(f"Encoded {len(HALLMARKS)} hallmarks.")
+    except Exception as e:
+        print(f"ERROR: Failed to BPE encode or convert hallmarks for mask: {e}")
+        raise
+
+    print("Encoding separators for mask...")
+    SEPS = [", ", ",", "; ", ";", "|"]
+    try:
+        sep_bpes = bpe_encode_lines(SEPS, desc="BPE Separators (for mask)")
+        for bpe_list in sep_bpes:
+             ids, _ = tokens_to_ids(bpe_list)
+             allowed.update(ids)
+        print(f"Encoded {len(SEPS)} separators.")
+    except Exception as e:
+         print(f"ERROR: Failed to BPE encode or convert separators for mask: {e}")
+         raise
+
+    # Add EOS token - Check if eos_id is valid
+    if not isinstance(eos_id, int) or eos_id < 0 or eos_id >= vocab_size:
+         print(f"Warning: Invalid EOS ID ({eos_id}). Defaulting mask EOS to 0.")
+         effective_eos_id = 0
+    else:
+         effective_eos_id = eos_id
+    allowed.add(effective_eos_id)
+    print(f"Total allowed token IDs (including EOS {effective_eos_id}): {len(allowed)}")
+
+    # Create the mask tensor on CPU first
+    mask = torch.full((vocab_size,), float('-inf'), device=torch.device('cpu'))
+    try:
+        # Filter out potential invalid IDs before creating list for indexing
+        # Ensure pad_id is valid if used for filtering
+        effective_pad_id = pad_id if isinstance(pad_id, int) and 0 <= pad_id < vocab_size else -1 # Use -1 if pad_id is invalid
+
+        valid_allowed_ids = []
+        for id_ in allowed:
+            if isinstance(id_, int) and 0 <= id_ < vocab_size: # Check type and range
+                 # Filter out pad_id unless it's the same as the effective_eos_id
+                 if id_ != effective_pad_id or id_ == effective_eos_id:
+                      valid_allowed_ids.append(id_)
+            # else: print(f"Warning: Invalid ID {id_} in allowed set skipped.") # Reduce noise
+
+        if not valid_allowed_ids:
+            raise ValueError("No valid token IDs found to allow in the mask.")
+
+        # Check ranges again after filtering (belt and braces)
+        max_valid_id = max(valid_allowed_ids)
+        min_valid_id = min(valid_allowed_ids)
+        if max_valid_id >= vocab_size or min_valid_id < 0:
+             # This should ideally not happen if filtering worked
+             raise IndexError(f"Filtered allowed IDs still out of range [{min_valid_id}, {max_valid_id}] for vocab size {vocab_size}.")
+
+        # Apply mask
+        mask[valid_allowed_ids] = 0.0 # Use list directly
+        print(f"Mask created with {len(valid_allowed_ids)} allowed indices.")
+
+    except IndexError as e:
+        print(f"ERROR: Index error while creating mask. Vocab size: {vocab_size}. Error: {e}")
+        # Find problematic IDs more carefully
+        problem_ids = [i for i in allowed if not isinstance(i, int) or i < 0 or i >= vocab_size]
+        print(f"Problematic IDs in allowed set: {problem_ids}")
+        raise
+    except Exception as e:
+        print(f"ERROR: Unexpected error creating mask: {e}")
+        traceback.print_exc()
+        raise
+
+    # Move final mask to target device
+    try:
+        target_device = torch.device(device) # Ensure device is a torch.device object
+        return mask.to(target_device)
+    except Exception as e:
+        print(f"Error moving mask to device '{device}': {e}")
+        raise
+
+
+# --- Global Variables for Loaded Model and Assets ---
+inference_model = None
+ALLOWED_MASK = None
+model_device = "cpu"
+config = None # Added global config
+
+# --- Initialization Function ---
+def initialize_model_and_tokenizer():
+    global inference_model, ALLOWED_MASK, model_device, token2id, config # Add config
+
+    print("Initializing model...")
+    # Determine device
+    model_device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device: {model_device}")
+
+    # Load tokenizer data first (essential for vocab size)
+    try:
+        load_tokenizer_data(DICT_TXT_PATH)
+        if not token2id: # Check if loading actually populated the dict
+             raise ValueError("Tokenizer loading failed to populate token2id dictionary.")
+    except Exception as e:
+        print(f"FATAL: Could not load tokenizer data. Cannot proceed. Error: {e}")
+        return False # Indicate failure
+
+    # Define model config (MUST match finetuning config)
+    try:
+        # Ensure config is globally accessible after definition
+        config = GPTConfig(
+            vocab_size=len(token2id), # Get vocab size from loaded data
+            block_size=128,           # Match training
+            n_layer=6,                # Match training
+            n_head=6,                 # Match training
+            n_embd=384,               # Match training
+            dropout=0.1,              # Match training (dropout is off in eval mode)
+            bias=True                 # Match training
+        )
+        print(f"Model Config: {config}")
+    except Exception as e:
+        print(f"FATAL: Error creating GPTConfig: {e}")
+        return False
+
+    # Instantiate base and wrapped model (on CPU initially)
+    try:
+        base_gpt = GPT(config)
+        inference_model = GPTWithSoftPrompt(base_gpt, prompt_len=1)
+    except Exception as e:
+        print(f"FATAL: Error instantiating model: {e}")
+        traceback.print_exc()
+        return False
+
+    # Load finetuned weights
+    print(f"Loading finetuned weights from: {FINETUNED_MODEL_PATH}")
+    if not os.path.exists(FINETUNED_MODEL_PATH):
+        print(f"ERROR: Model weights file not found at {FINETUNED_MODEL_PATH}")
+        return False
+
+    try:
+        # Load state dict onto CPU first
+        state_dict = torch.load(FINETUNED_MODEL_PATH, map_location='cpu')
+
+        # Clean state dict keys (handle DDP 'module.' prefix)
+        cleaned_state_dict = {}
+        for k, v in state_dict.items():
+            name = k[7:] if k.startswith('module.') else k
+            cleaned_state_dict[name] = v
+
+        # Load into model
+        missing_keys, unexpected_keys = inference_model.load_state_dict(cleaned_state_dict, strict=False)
+        if missing_keys:
+            # Filter out non-persistent buffer keys if necessary (though strict=False should handle this)
+            missing_persistent = [k for k in missing_keys if inference_model.get_parameter(k) is not None or inference_model.get_buffer(k) is not None]
+            if missing_persistent:
+                print("Warning: Missing persistent keys during state dict load:", missing_persistent)
+        if unexpected_keys:
+            print("Warning: Unexpected keys during state dict load:", unexpected_keys)
+        print("Weights loaded successfully.")
+
+    except Exception as e:
+        print(f"Error loading state dict from {FINETUNED_MODEL_PATH}: {e}")
+        print("Ensure the model architecture matches the saved checkpoint and the file is not corrupted.")
+        traceback.print_exc()
+        return False
+
+    # Move model to target device and set to eval mode
+    try:
+        inference_model.to(model_device)
+        inference_model.eval()
+        print(f"Model moved to device: {model_device} and set to eval mode.")
+    except Exception as e:
+        print(f"Error moving model to device '{model_device}': {e}")
+        traceback.print_exc()
+        return False
+
+
+    # Build the allowed token mask (after model is on device)
+    print("Building allowed token mask...")
+    try:
+        if config.vocab_size <= 0:
+             raise ValueError("Vocabulary size must be positive to build mask.")
+        # Ensure model_device is valid before passing
+        device_obj = torch.device(model_device)
+        ALLOWED_MASK = build_allowed_token_mask(config.vocab_size, device_obj)
+        print("Allowed token mask created.")
+    except Exception as e:
+        print(f"ERROR: Failed to build allowed token mask: {e}")
+        traceback.print_exc()
+        return False
+
+    return True # Indicate success
+
+
+# --- Inference Function ---
+def predict_hallmarks(abstract_text):
+    global inference_model, ALLOWED_MASK, model_device, token2id, eos_id
+
+    # --- Pre-computation Checks ---
+    if inference_model is None:
+         print("Error: Inference model is not loaded.")
+         return ["Error: Model not loaded"]
+    if ALLOWED_MASK is None:
+         print("Error: Allowed mask is not built.")
+         return ["Error: Mask not built"]
+    if not token2id:
+         print("Error: Tokenizer vocabulary not loaded.")
+         return ["Error: Tokenizer not loaded"]
+
+    # --- Input Validation ---
+    if not isinstance(abstract_text, str):
+        print(f"Warning: Received non-string abstract text type: {type(abstract_text)}. Attempting conversion.")
+        try:
+            abstract_text = str(abstract_text)
+        except Exception:
+             return ["Error: Invalid input type"]
+    if not abstract_text.strip():
+        print("Warning: Received empty or whitespace-only abstract text.")
+        return [] # Return empty list for empty input
+
+
+    try:
+        # --- 1. Preprocess and Tokenize Input ---
+        print("Tokenizing input abstract...")
+        cleaned_abstract = " ".join(abstract_text.split())
+        if not cleaned_abstract:
+             print("Warning: Input abstract contains only whitespace after cleaning.")
+             return []
+
+        bpe_tokens_list = bpe_encode_lines([cleaned_abstract])
+        if not bpe_tokens_list or not bpe_tokens_list[0]: # Check if list or first element is empty
+            print("Warning: BPE encoding resulted in empty tokens.")
+            return []
+        bpe_tokens = bpe_tokens_list[0]
+
+        input_ids_list, oov = tokens_to_ids(bpe_tokens)
+        if oov > 0:
+             print(f"Info: Input contained {oov} OOV tokens.")
+
+        # Add EOS token
+        input_ids = input_ids_list + [eos_id]
+
+        # Convert to tensor and move to device
+        input_tensor = torch.tensor(input_ids, dtype=torch.long).unsqueeze(0).to(model_device)
+
+        # --- 2. Generate Predictions ---
+        print("Generating predictions...")
+        with torch.no_grad():
+            generated_ids_tensor = inference_model.generate_labels(
+                input_tensor,
+                allowed_mask=ALLOWED_MASK,
+                max_new_tokens=30,
+                temperature=0.0
+            )
+
+        # --- 3. Decode and Post-process ---
+        print("Decoding and cleaning predictions...")
+        if generated_ids_tensor is None or generated_ids_tensor.numel() == 0:
+             print("Warning: Generation resulted in empty tensor.")
+             generated_ids = []
+        else:
+             # Ensure tensor is on CPU before converting to list
+             generated_ids = generated_ids_tensor[0].cpu().tolist()
+
+        if not generated_ids:
+             print("No tokens generated.")
+             return []
+
+        generated_tokens = ids_to_tokens(generated_ids)
+
+        # Remove tokens after EOS if present
+        try:
+            eos_token_str = id2token.get(eos_id, "</s>") # Get string representation
+            if eos_token_str in generated_tokens:
+                 eos_index = generated_tokens.index(eos_token_str)
+                 generated_tokens = generated_tokens[:eos_index]
+        except ValueError:
+             pass # EOS not found is okay
+
+        # Decode BPE tokens to string
+        generated_text = bpe_decode_tokens(generated_tokens).strip().lower()
+        print(f"Raw generated text: '{generated_text}'")
+
+        # Split potential multi-labels and map to canonical
+        parts = []
+        if generated_text:
+             potential_parts = re.split(r'[;,|]\s*', generated_text)
+             parts = [p.strip() for p in potential_parts if p.strip()]
+             if not parts: # Handle case with no delimiters
+                  parts = [generated_text]
+
+        predicted_labels = []
+        seen_labels = set()
+        for p in parts:
+            canonical_label = to_canonical(p)
+            if canonical_label and canonical_label not in seen_labels:
+                predicted_labels.append(canonical_label)
+                seen_labels.add(canonical_label)
+
+        print(f"Final predicted labels: {predicted_labels}")
+        return predicted_labels
+
+    # --- Error Handling ---
+    except FileNotFoundError as fnf_err:
+         print(f"ERROR during prediction (File Not Found - likely BPE related): {fnf_err}")
+         traceback.print_exc()
+         return ["Error: BPE file processing error"]
+    except PermissionError as perm_err:
+        print(f"ERROR during prediction (Permission Error - likely fastBPE): {perm_err}")
+        traceback.print_exc()
+        return ["Error: BPE execution permission"]
+    except RuntimeError as run_err:
+         if "CUDA out of memory" in str(run_err):
+              print(f"ERROR: CUDA Out of Memory during prediction. Input length: {len(input_ids) if 'input_ids' in locals() else 'N/A'}")
+              traceback.print_exc()
+              return ["Error: Input too long (OOM)"]
+         else:
+              print(f"ERROR during prediction (PyTorch RuntimeError): {run_err}")
+              traceback.print_exc()
+              return ["Error: Model runtime error"]
+    except Exception as e:
+        print(f"ERROR during prediction (General Exception): {e}")
+        traceback.print_exc()
+        return [f"Error: An unexpected error occurred"]
+
+
+# --- Flask App ---
+app = Flask(__name__)
+
+# --- Load Model on Startup ---
+model_initialized = False
+
+@app.before_request
+def ensure_model_loaded():
+    """ Ensures model is loaded before handling the first request. """
+    global model_initialized
+    if not model_initialized:
+        print("First request received, attempting to initialize model...")
+        # Add basic locking if deploying with multiple workers (though not fully thread-safe here)
+        # For true multi-worker safety, model loading should happen before workers fork.
+        try:
+            if initialize_model_and_tokenizer():
+                model_initialized = True
+                print("Model initialization successful.")
+            else:
+                print("FATAL: Model initialization failed during first request.")
+                # We won't raise an error here, but subsequent requests will fail until fixed.
+        except Exception as init_err:
+             print(f"FATAL: Exception during model initialization: {init_err}")
+             traceback.print_exc()
+
+
+# --- Routes ---
+@app.route('/')
+def home():
+    """ Renders the HTML frontend page. """
+    # Check if initialization failed and show an error page if so?
+    # For simplicity, we assume initialization works or subsequent predict calls fail.
+    return render_template('index.html')
+
+@app.route('/predict', methods=['POST'])
+def predict():
+    """ Handles prediction requests from the frontend. """
+    global model_initialized
+    # Check if model is ready
+    if not model_initialized:
+         print("Error: Model not initialized when /predict called.")
+         # Return a specific status code like Service Unavailable
+         return jsonify({'error': 'Model is not ready. Please try again later.'}), 503
+
+    # Validate request format
+    if not request.is_json:
+        return jsonify({'error': 'Request must be JSON'}), 400
+
+    data = request.get_json()
+    abstract = data.get('abstract')
+
+    # Validate input abstract
+    if not abstract:
+        return jsonify({'error': 'Missing "abstract" field in JSON request'}), 400
+    if not isinstance(abstract, str):
+         return jsonify({'error': '"abstract" field must be a string'}), 400
+    if len(abstract.strip()) == 0:
+        print("Received empty abstract, returning empty prediction.")
+        return jsonify({'predictions': []})
+    MAX_ABSTRACT_LEN = 10000 # Define max length
+    if len(abstract) > MAX_ABSTRACT_LEN:
+         print(f"Received overly long abstract ({len(abstract)} chars), rejecting.")
+         return jsonify({'error': f'Input abstract is too long (max {MAX_ABSTRACT_LEN} chars)'}), 413 # Payload Too Large
+
+    print(f"\n--- Received Prediction Request ---")
+    print(f"Input Abstract (first 100 chars): {abstract[:100]}...")
+
+    try:
+        # Perform prediction
+        predictions = predict_hallmarks(abstract)
+        print(f"--- Prediction Complete ---")
+
+        # Check if the result indicates an internal error occurred
+        if isinstance(predictions, list) and len(predictions) > 0 and predictions[0].startswith("Error:"):
+             print(f"Internal error during prediction: {predictions[0]}")
+             # Return a generic server error to the client
+             return jsonify({'error': 'An internal error occurred during prediction.'}), 500
+        else:
+             # Return successful predictions
+             return jsonify({'predictions': predictions})
+
+    except Exception as e:
+        # Catch unexpected errors in the route handler itself
+        print(f"--- Prediction Failed Unexpectedly in Route ---")
+        print(f"Error: {e}")
+        traceback.print_exc()
+        return jsonify({'error': 'An internal server error occurred.'}), 500
+
+# --- Run the App ---
+if __name__ == '__main__':
+    # Initialize model eagerly when running script directly
+    if not model_initialized:
+        print("Running script directly, initializing model eagerly...")
+        if initialize_model_and_tokenizer():
+            model_initialized = True
+            print("Model initialization successful.")
+        else:
+            print("FATAL: Model initialization failed. Cannot start Flask server.")
+            exit(1) # Exit if model fails to load on startup
+
+    # Check fastBPE path validity before starting server
+    abs_fastbpe_path = os.path.abspath(FASTBPE_BIN_PATH)
+    if not os.path.exists(abs_fastbpe_path):
+         print(f"ERROR: fastBPE binary not found at '{abs_fastbpe_path}'.")
+         print("Please ensure fastBPE is compiled and the path is correct relative to app.py.")
+         exit(1)
+    if not os.access(abs_fastbpe_path, os.X_OK):
+        print(f"ERROR: fastBPE binary at '{abs_fastbpe_path}' is not executable.")
+        print("Attempting to make it executable with 'chmod +x'...")
+        try:
+            os.chmod(abs_fastbpe_path, 0o755)
+            print(f"Successfully made '{abs_fastbpe_path}' executable.")
+        except OSError as e:
+            print(f"ERROR: Failed to make fastBPE executable: {e}")
+            print("Please set execute permissions manually (e.g., 'chmod +x ./fast').")
+            exit(1)
+
+    print("Starting Flask server...")
+    # Use host='0.0.0.0' to make it accessible on your network
+    # Set debug=False for production environments
+    app.run(host='0.0.0.0', port=5000, debug=False) # Changed debug to False
+

fast

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c86a282547b0bc4d088adee84b45ad031db1b23139d944a159ae8f9cfce3186e
+size 132336

hoc_best.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:13d02e9a35bd2fc935f5f1291a16b568928e223cba9152f5284ef45755f91491
+size 107913359

requirements.txt

@@ -0,0 +1,7 @@
+Flask
+torch
+sacremoses==0.0.53
+numpy
+pandas
+tqdm
+waitress

templates/index.html

@@ -0,0 +1,191 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>BioGPT HoC Inference</title>
+    <!-- Load Tailwind CSS -->
+    <script src="https://cdn.tailwindcss.com"></script>
+    <!-- Include Inter font -->
+    <link rel="preconnect" href="https://fonts.googleapis.com">
+    <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
+    <link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap" rel="stylesheet">
+    <style>
+        body {
+            font-family: 'Inter', sans-serif;
+        }
+        .hallmark-badge {
+            display: inline-block;
+            padding: 0.3rem 0.6rem;
+            margin: 0.2rem;
+            border-radius: 0.5rem; /* Rounded corners */
+            font-size: 0.8rem;
+            line-height: 1rem;
+            font-weight: 500;
+            cursor: default;
+            transition: background-color 0.2s ease-in-out;
+            border: 1px solid transparent;
+        }
+        /* Color definitions (same as before) */
+        .hallmark-color-0 { background-color: #EFF6FF; color: #1E40AF; border-color: #BEE3F8; } /* Blue */
+        .hallmark-color-1 { background-color: #ECFDF5; color: #065F46; border-color: #A7F3D0; } /* Green */
+        .hallmark-color-2 { background-color: #FFFBEB; color: #92400E; border-color: #FDE68A; } /* Yellow */
+        .hallmark-color-3 { background-color: #FEF2F2; color: #991B1B; border-color: #FECACA; } /* Red */
+        .hallmark-color-4 { background-color: #F5F3FF; color: #5B21B6; border-color: #DDD6FE; } /* Purple */
+        .hallmark-color-5 { background-color: #FDF2F8; color: #9D174D; border-color: #FBCFE8; } /* Pink */
+        .hallmark-color-6 { background-color: #EEF2FF; color: #3730A3; border-color: #C7D2FE; } /* Indigo */
+        .hallmark-color-7 { background-color: #F0FDFA; color: #134E4A; border-color: #99F6E4; } /* Teal */
+        .hallmark-color-8 { background-color: #FFF7ED; color: #9A3412; border-color: #FED7AA; } /* Orange */
+        .hallmark-color-9 { background-color: #F9FAFB; color: #374151; border-color: #E5E7EB; } /* Gray */
+
+        /* Loading Spinner */
+        .spinner {
+            border: 4px solid rgba(0, 0, 0, 0.1);
+            width: 36px;
+            height: 36px;
+            border-radius: 50%;
+            border-left-color: #4f46e5; /* Indigo */
+            animation: spin 1s ease infinite;
+            margin: 20px auto;
+        }
+        @keyframes spin {
+            0% { transform: rotate(0deg); }
+            100% { transform: rotate(360deg); }
+        }
+    </style>
+</head>
+<body class="bg-gray-100 min-h-screen flex flex-col items-center justify-center p-4 sm:p-6">
+
+    <div class="bg-white p-6 sm:p-8 rounded-lg shadow-xl w-full max-w-2xl space-y-6">
+
+        <!-- Header -->
+        <div class="text-center border-b pb-4">
+            <h1 class="text-2xl sm:text-3xl font-bold text-gray-800 mb-1">BioGPT HoC Inference</h1>
+            <h2 class="text-md sm:text-lg font-semibold text-blue-700 mb-2">~27M Parameter Finetuned SLM</h2>
+            <p class="text-sm text-gray-600 max-w-2xl mx-auto">
+                Enter a biomedical abstract below to predict Hallmarks of Cancer labels using the finetuned model.
+            </p>
+        </div>
+
+        <!-- Input Area -->
+        <div>
+            <label for="abstract-input" class="block text-sm font-medium text-gray-700 mb-1">Enter Abstract Text:</label>
+            <textarea id="abstract-input" rows="8" class="w-full p-3 border border-gray-300 rounded-lg focus:ring-2 focus:ring-blue-500 focus:border-transparent transition duration-150 ease-in-out shadow-sm" placeholder="Paste or type abstract here..."></textarea>
+        </div>
+
+        <!-- Predict Button -->
+        <div class="text-center">
+            <button id="predict-button" class="bg-blue-600 hover:bg-blue-700 text-white font-semibold py-2.5 px-8 rounded-lg transition duration-150 ease-in-out shadow-md hover:shadow-lg focus:outline-none focus:ring-2 focus:ring-blue-500 focus:ring-offset-2 disabled:opacity-50 disabled:cursor-not-allowed">
+                Predict Hallmarks
+            </button>
+        </div>
+
+        <!-- Results Area -->
+        <div id="results-area" class="space-y-3 hidden">
+             <div class="border border-green-200 p-4 rounded-lg bg-green-50 shadow-sm min-h-[80px]">
+                 <h3 class="text-md font-semibold text-green-800 mb-2">Predicted Hallmarks:</h3>
+                 <div id="loading-indicator" class="hidden text-center">
+                     <div class="spinner"></div>
+                     <p class="text-sm text-gray-600">Predicting...</p>
+                 </div>
+                 <div id="error-message" class="hidden text-center text-red-600 font-medium">
+                     An error occurred. Please try again.
+                 </div>
+                 <div id="prediction-output">
+                     <!-- Badges will be inserted here -->
+                 </div>
+            </div>
+        </div>
+
+    </div>
+
+    <script>
+        // --- Hallmark Colors (for consistent styling) ---
+        const hallmarkList = [ // Ensure this matches the list in app.py
+            "activating invasion and metastasis", "avoiding immune destruction",
+            "cellular energetics", "enabling replicative immortality",
+            "evading growth suppressors", "genomic instability and mutation",
+            "inducing angiogenesis", "resisting cell death",
+            "sustaining proliferative signaling", "tumor promoting inflammation",
+        ];
+        const hallmarkColors = {};
+        hallmarkList.forEach((hallmark, index) => {
+            hallmarkColors[hallmark] = `hallmark-color-${index % 10}`;
+        });
+
+        // --- DOM Elements ---
+        const abstractInput = document.getElementById('abstract-input');
+        const predictButton = document.getElementById('predict-button');
+        const resultsArea = document.getElementById('results-area');
+        const loadingIndicator = document.getElementById('loading-indicator');
+        const errorMessage = document.getElementById('error-message');
+        const predictionOutput = document.getElementById('prediction-output');
+
+        // --- Event Listener ---
+        predictButton.addEventListener('click', handlePrediction);
+
+        // --- Prediction Logic ---
+        async function handlePrediction() {
+            const abstractText = abstractInput.value.trim();
+            if (!abstractText) {
+                alert("Please enter some abstract text.");
+                return;
+            }
+
+            // --- UI Updates: Start loading ---
+            predictButton.disabled = true;
+            resultsArea.classList.remove('hidden');
+            predictionOutput.innerHTML = ''; // Clear previous results
+            errorMessage.classList.add('hidden');
+            loadingIndicator.classList.remove('hidden');
+
+            try {
+                // --- Call the backend API ---
+                const response = await fetch('/predict', {
+                    method: 'POST',
+                    headers: {
+                        'Content-Type': 'application/json',
+                    },
+                    body: JSON.stringify({ abstract: abstractText }),
+                });
+
+                // --- Handle Response ---
+                loadingIndicator.classList.add('hidden'); // Hide loading indicator
+                if (!response.ok) {
+                    const errorData = await response.json();
+                    throw new Error(errorData.error || `HTTP error! status: ${response.status}`);
+                }
+
+                const data = await response.json();
+                displayPredictions(data.predictions);
+
+            } catch (error) {
+                console.error('Prediction failed:', error);
+                loadingIndicator.classList.add('hidden');
+                errorMessage.textContent = `Prediction failed: ${error.message}`;
+                errorMessage.classList.remove('hidden');
+            } finally {
+                 predictButton.disabled = false; // Re-enable button
+            }
+        }
+
+        // --- Display Logic ---
+        function displayPredictions(predictions) {
+            predictionOutput.innerHTML = ''; // Clear previous just in case
+            errorMessage.classList.add('hidden');
+
+            if (predictions && predictions.length > 0) {
+                predictions.forEach(label => {
+                    const badge = document.createElement('span');
+                    badge.textContent = label;
+                    const colorClass = hallmarkColors[label] || 'hallmark-color-9'; // Default color
+                    badge.className = `hallmark-badge ${colorClass}`;
+                    predictionOutput.appendChild(badge);
+                });
+            } else {
+                predictionOutput.innerHTML = '<span class="text-gray-500 italic text-sm">No specific hallmarks predicted for this abstract.</span>';
+            }
+        }
+    </script>
+</body>
+</html>