adding finetune code

bio_gpt_finetune.py

@@ -0,0 +1,2082 @@
+# -*- coding: utf-8 -*-
+"""Vizuara BioGPT from Scratch.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/1ys-b99GalAtTE9m7bGwCCACZYv2M8HjO
+
+#Vizuara AI Labs: BioGPT Pre-training + Finetuning
+
+## Part 1: Pre-training
+
+### 1.1 Loading the dataset
+"""
+
+# Colab: Download ~10 GB (uncompressed) of PubMed baseline XML
+import os, re, subprocess, math, requests
+from bs4 import BeautifulSoup
+from urllib.parse import urljoin
+
+BASE_URL = "https://ftp.ncbi.nlm.nih.gov/pubmed/baseline/"
+TARGET_UNCOMPRESSED_GB = 1.0
+DEST = "/content/pubmed_xml_subset"
+os.makedirs(DEST, exist_ok=True)
+
+# 1) Fetch list of .gz files from the baseline index
+html = requests.get(BASE_URL, timeout=60).text
+soup = BeautifulSoup(html, "html.parser")
+
+# All .gz files (e.g., pubmed24n0001.xml.gz)
+hrefs = [a.get("href") for a in soup.find_all("a", href=True)]
+gz_files = sorted([h for h in hrefs if h.endswith(".gz")])
+
+print(f"Found {len(gz_files)} .gz files on the baseline index.")
+
+# 2) Download sequentially until uncompressed total ≈ target
+def gz_uncompressed_bytes(local_path):
+    # Use gzip -l to read uncompressed size from footer (fast; no full decompress)
+    out = subprocess.check_output(["gzip", "-l", local_path]).decode()
+    # The second line has: compressed  uncompressed  ratio  uncompressed_name
+    lines = out.strip().splitlines()
+    if len(lines) >= 2:
+        parts = re.split(r"\s+", lines[1].strip())
+        # parts[1] = uncompressed bytes
+        return int(parts[1])
+    return 0
+
+total_uncompressed = 0
+downloaded = []
+
+for fname in gz_files:
+    url = urljoin(BASE_URL, fname)
+    local = os.path.join(DEST, fname)
+    if not os.path.exists(local):
+        print(f"→ downloading {fname} ...")
+        # quiet, continue on partial, retry a bit
+        ret = subprocess.call(["wget", "-q", "-c", "-O", local, url])
+        if ret != 0:
+            print(f"  ! failed: {fname}; skipping")
+            if os.path.exists(local): os.remove(local)
+            continue
+    # read uncompressed size
+    try:
+        ub = gz_uncompressed_bytes(local)
+        total_uncompressed += ub
+        downloaded.append((fname, ub))
+        print(f"  added {fname}: {ub/1e9:.3f} GB uncompressed  |  total ≈ {total_uncompressed/1e9:.3f} GB")
+    except Exception as e:
+        print(f"  ! could not read size for {fname}: {e}")
+
+    if total_uncompressed >= TARGET_UNCOMPRESSED_GB * 1e9:
+        print("\nTarget reached. Stopping downloads.")
+        break
+
+print(f"\nDone. Saved {len(downloaded)} files to: {DEST}")
+print(f"Approx. uncompressed total: {total_uncompressed/1e9:.3f} GB")
+
+"""### 1.2 Converting title and abstract from XML to TXT"""
+
+# Colab cell: Parse title + abstract to plain text (one doc/line)
+import os, gzip, glob
+from lxml import etree
+from tqdm import tqdm
+
+SRC_DIR = "/content/pubmed_xml_subset"         # where your .xml.gz files are
+OUT_DIR = "/content/pubmed_txt"                # output folder
+os.makedirs(OUT_DIR, exist_ok=True)
+
+train_path = f"{OUT_DIR}/train.txt"
+valid_path = f"{OUT_DIR}/valid.txt"
+test_path  = f"{OUT_DIR}/test.txt"
+
+# ----- helper: stream-parse one PubMed file -----
+def yield_title_abstract(fp):
+    # iterparse to avoid loading whole XML into RAM
+    ctx = etree.iterparse(gzip.open(fp), events=("end",), tag="PubmedArticle")
+    for _, elem in ctx:
+        # Title
+        t = elem.find(".//ArticleTitle")
+        title = (t.text or "").strip() if t is not None else ""
+        # Abstract may have multiple parts <AbstractText>
+        abs_nodes = elem.findall(".//AbstractText")
+        abs_parts = []
+        for a in abs_nodes:
+            txt = (a.text or "").strip()
+            if txt:
+                abs_parts.append(txt)
+        abstract = " ".join(abs_parts).strip()
+
+        if title and abstract:
+            text = f"{title}. {abstract}"
+            # clean newlines/tabs
+            text = " ".join(text.split())
+            yield text
+
+        # free memory
+        elem.clear()
+        while elem.getprevious() is not None:
+            del elem.getparent()[0]
+    del ctx
+
+# ----- collect and write -----
+gz_files = sorted(glob.glob(os.path.join(SRC_DIR, "*.xml.gz")))
+print(f"Found {len(gz_files)} gz files")
+
+# We'll stream all docs, then do a simple split by count.
+all_out = f"{OUT_DIR}/_all.txt"
+with open(all_out, "w", encoding="utf-8") as out:
+    for fp in tqdm(gz_files, desc="Parsing"):
+        for line in yield_title_abstract(fp):
+            out.write(line + "\n")
+
+# Quick stats
+num_lines = sum(1 for _ in open(all_out, "r", encoding="utf-8"))
+print("Total docs with title+abstract:", num_lines)
+
+# Split 98% / 1% / 1% (adjust if you like)
+train_n = int(num_lines * 0.98)
+valid_n = int(num_lines * 0.01)
+test_n  = num_lines - train_n - valid_n
+
+with open(all_out, "r", encoding="utf-8") as fin, \
+     open(train_path, "w", encoding="utf-8") as ftr, \
+     open(valid_path, "w", encoding="utf-8") as fva, \
+     open(test_path,  "w", encoding="utf-8") as fte:
+    for i, line in enumerate(fin):
+        if i < train_n: ftr.write(line)
+        elif i < train_n + valid_n: fva.write(line)
+        else: fte.write(line)
+
+print("Wrote:")
+print(" ", train_path)
+print(" ", valid_path)
+print(" ", test_path)
+
+# Commented out IPython magic to ensure Python compatibility.
+# Colab cell: Install tools
+!pip -q install sacremoses==0.0.53
+!sudo apt-get -y install g++ >/dev/null
+
+# fastBPE (build once)
+!git clone -q https://github.com/glample/fastBPE.git /content/fastBPE
+# %cd /content/fastBPE
+!g++ -std=c++11 -O3 -pthread fastBPE/main.cc -IfastBPE -o fast
+# %cd /content
+
+# fairseq (0.12.0 recommended for GPT2-medium arch flag)
+!git clone -q https://github.com/pytorch/fairseq.git /content/fairseq
+# %cd /content/fairseq
+!git checkout v0.12.0 -q
+!pip -q install .
+# %cd /content
+
+"""### 1.3 Fetch the BioGPT Vocabulary and merged tokens"""
+
+# Colab cell: Grab BioGPT bpecodes/dict
+!wget -q -O /content/bpecodes https://raw.githubusercontent.com/microsoft/BioGPT/main/data/BioGPT/bpecodes
+!wget -q -O /content/dict.txt  https://raw.githubusercontent.com/microsoft/BioGPT/main/data/BioGPT/dict.txt
+!wc -l /content/dict.txt && head -n 5 /content/dict.txt
+
+"""### 1.4 Use Moses tokenizer to clean text before applying BPE"""
+
+import os
+from sacremoses import MosesTokenizer
+from tqdm.auto import tqdm
+
+TXT_DIR = "/content/pubmed_txt"
+BPE_DIR = "/content/pubmed_bpe"
+os.makedirs(BPE_DIR, exist_ok=True)
+
+mt = MosesTokenizer(lang="en")
+
+def tokenize_file(in_path, out_path, show_progress=True):
+    # Count lines once for a nice total
+    with open(in_path, "r", encoding="utf-8") as f:
+        total = sum(1 for _ in f)
+
+    with open(in_path, "r", encoding="utf-8") as fin, \
+         open(out_path, "w", encoding="utf-8") as fout:
+        iterator = fin
+        if show_progress:
+            iterator = tqdm(fin, total=total, desc=f"Tokenizing {os.path.basename(in_path)}")
+        for line in iterator:
+            line = line.strip()
+            if not line:
+                continue
+            fout.write(mt.tokenize(line, return_str=True) + "\n")
+
+for split in ["train", "valid", "test"]:
+    tok = f"{BPE_DIR}/{split}.tok"
+    bpe = f"{BPE_DIR}/{split}.bpe"
+    tokenize_file(f"{TXT_DIR}/{split}.txt", tok)
+
+"""### 1.5 Apply BPE to dataset"""
+
+# Commented out IPython magic to ensure Python compatibility.
+import os, math, subprocess, numpy as np, shutil
+from tqdm.auto import tqdm
+
+BPE_CODES = "/content/bpecodes"     # BioGPT bpecodes
+DICT_TXT  = "/content/dict.txt"     # BioGPT dict
+BPE_DIR   = "/content/pubmed_bpe"   # where your .tok files are
+BIN_DIR   = "/content/pubmed_memmap"
+TMP_DIR   = "/content/_bpe_tmp"
+os.makedirs(BIN_DIR, exist_ok=True)
+os.makedirs(TMP_DIR, exist_ok=True)
+
+# --- load vocab ---
+token2id = {}
+with open(DICT_TXT, encoding="utf-8") as f:
+    for i, line in enumerate(f):
+        tok = line.split()[0]
+        token2id[tok] = i
+# choose a fallback id ONLY IF we see OOVs later
+fallback_id = token2id.get("</s>", next(iter(token2id.values())))  # prefer EOS, else first token
+
+# --- ensure fastBPE binary exists ---
+if not os.path.exists("/content/fastBPE/fast"):
+    !git clone -q https://github.com/glample/fastBPE.git /content/fastBPE
+#     %cd /content/fastBPE
+    !g++ -std=c++11 -O3 -pthread fastBPE/main.cc -IfastBPE -o fast
+#     %cd /content
+
+def line_count(path):
+    c = 0
+    with open(path, encoding="utf-8") as f:
+        for _ in f:
+            c += 1
+    return c
+
+def apply_bpe_with_progress(tok_file, bpe_file, shards=50):
+    total_lines = line_count(tok_file)
+    if total_lines == 0:
+        open(bpe_file, "w").close()
+        return
+
+    shards = max(1, min(shards, total_lines))
+    lines_per = math.ceil(total_lines / shards)
+
+    split_dir = os.path.join(TMP_DIR, "split")
+    out_dir   = os.path.join(TMP_DIR, "bpe_parts")
+    os.makedirs(split_dir, exist_ok=True)
+    os.makedirs(out_dir, exist_ok=True)
+
+    # 1) split with progress
+    with open(tok_file, encoding="utf-8") as fin:
+        shard_idx = 0
+        line_idx  = 0
+        fout = None
+        pbar = tqdm(total=total_lines, desc=f"Splitting {os.path.basename(tok_file)}")
+        for line in fin:
+            if line_idx % lines_per == 0:
+                if fout: fout.close()
+                shard_idx += 1
+                fout = open(os.path.join(split_dir, f"part_{shard_idx:05d}.tok"), "w", encoding="utf-8")
+            fout.write(line)
+            line_idx += 1
+            pbar.update(1)
+        if fout: fout.close()
+        pbar.close()
+
+    # 2) BPE on each shard with progress
+    parts = sorted([p for p in os.listdir(split_dir) if p.endswith(".tok")])
+    for p in tqdm(parts, desc="Applying BPE to shards"):
+        src = os.path.join(split_dir, p)
+        dst = os.path.join(out_dir, p.replace(".tok", ".bpe"))
+        subprocess.check_call(["/content/fastBPE/fast", "applybpe", dst, src, BPE_CODES])
+
+    # 3) concat with progress
+    with open(bpe_file, "w", encoding="utf-8") as fout:
+        for p in tqdm(parts, desc="Concatenating BPE shards"):
+            src = os.path.join(out_dir, p.replace(".tok", ".bpe"))
+            with open(src, encoding="utf-8") as fin:
+                shutil.copyfileobj(fin, fout)
+
+    shutil.rmtree(split_dir, ignore_errors=True)
+    shutil.rmtree(out_dir,   ignore_errors=True)
+
+def make_bin(split, dtype=np.uint16, shards=64):
+    tok_file = os.path.join(BPE_DIR, f"{split}.tok")
+    bpe_file = os.path.join(BPE_DIR, f"{split}.bpe")
+
+    print(f"\n[{split}] Step 1: Applying BPE merges with progress...")
+    apply_bpe_with_progress(tok_file, bpe_file, shards=shards)
+
+    print(f"[{split}] Step 2: Counting total tokens...")
+    total_tokens, total_lines = 0, 0
+    with open(bpe_file, encoding="utf-8") as f:
+        for line in tqdm(f, desc="Counting tokens"):
+            total_tokens += len(line.strip().split())
+            total_lines  += 1
+    print(f"[{split}] Total tokens: {total_tokens:,} | lines: {total_lines:,}")
+
+    print(f"[{split}] Step 3: Encoding to IDs & writing memmap...")
+    bin_path = os.path.join(BIN_DIR, f"{split}.bin")
+    arr = np.memmap(bin_path, dtype=dtype, mode="w+", shape=(total_tokens,))
+
+    idx = 0
+    oov_count = 0
+    oov_samples = {}
+    with open(bpe_file, encoding="utf-8") as f:
+        for line in tqdm(f, total=total_lines, desc=f"Encoding {split}"):
+            toks = line.strip().split()
+            ids  = []
+            for t in toks:
+                if t in token2id:
+                    ids.append(token2id[t])
+                else:
+                    oov_count += 1
+                    if len(oov_samples) < 10:
+                        oov_samples[t] = oov_samples.get(t, 0) + 1
+                    ids.append(fallback_id)  # safe fallback if any OOVs occur
+            n = len(ids)
+            arr[idx:idx+n] = np.fromiter(ids, dtype=dtype, count=n)
+            idx += n
+    arr.flush()
+
+    if oov_count == 0:
+        print(f"[{split}] ✅ Saved {bin_path} (no OOVs)")
+    else:
+        print(f"[{split}] ⚠️ Saved {bin_path} with {oov_count} OOV tokens mapped to id {fallback_id}.")
+        print("   First few OOV examples:", list(oov_samples.items()))
+
+for split in ["train", "valid", "test"]:
+    make_bin(split, dtype=np.uint16, shards=64)
+
+"""### 1.6 Create input-output pairs"""
+
+import os, numpy as np, torch
+
+BIN_ROOT = "/content/pubmed_memmap"  # where your .bin files are
+DTYPE = np.uint16                    # you saved with uint16
+
+def get_batch(split):
+    fname = "train.bin" if split == "train" else "valid.bin"
+    path = os.path.join(BIN_ROOT, fname)
+    data = np.memmap(path, dtype=DTYPE, mode='r')
+
+    ix = torch.randint(len(data) - block_size, (batch_size,))
+    x = torch.stack([torch.from_numpy(data[i:i+block_size].astype(np.int64)) for i in ix])
+    y = torch.stack([torch.from_numpy(data[i+1:i+1+block_size].astype(np.int64)) for i in ix])
+
+    if device_type == 'cuda':
+        x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
+    else:
+        x, y = x.to(device), y.to(device)
+    return x, y
+
+"""### 1.7 Define BioGPT architecture"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from dataclasses import dataclass
+import numpy as np
+from tqdm.auto import tqdm
+from contextlib import nullcontext
+import os
+
+class LayerNorm(nn.Module):
+    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):
+    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')
+        if not self.flash:
+            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
+                                       .view(1, 1, config.block_size, config.block_size))
+
+    def forward(self, x):
+        B, T, C = x.size()
+        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)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+
+        if self.flash:
+            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:
+            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+            att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            y = att @ v
+
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+class MLP(nn.Module):
+    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):
+        return self.dropout(self.c_proj(self.gelu(self.c_fc(x))))
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln1 = LayerNorm(config.n_embd, config.bias)
+        self.attn = CausalSelfAttention(config)
+        self.ln2 = LayerNorm(config.n_embd, 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
+
+@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 GPT(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        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, config.bias),
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.transformer.wte.weight = self.lm_head.weight  # weight tying
+
+        self.apply(self._init_weights)
+        for pn, p in self.named_parameters():
+            if pn.endswith('c_proj.weight'):
+                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):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None):
+        device = idx.device
+        b, t = idx.size()
+        assert t <= self.config.block_size
+        pos = torch.arange(0, t, dtype=torch.long, device=device)
+
+        tok_emb = self.transformer.wte(idx)
+        pos_emb = self.transformer.wpe(pos)
+        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:
+            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:
+            logits = self.lm_head(x[:, [-1], :])
+            return logits, None
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens, temperature=1.0, top_k=None):
+        """
+        Generate tokens given a conditioning sequence.
+        idx: Tensor of shape (B, T)
+        """
+        for _ in range(max_new_tokens):
+            idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / temperature
+            if top_k is not None:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = -float('Inf')
+            probs = F.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat((idx, idx_next), dim=1)
+        return idx
+
+vocab_size = sum(1 for _ in open("/content/dict.txt", encoding="utf-8"))
+print("Vocab size:", vocab_size)  # should be ~42380
+
+"""### 1.8 Define configuration"""
+
+# Pick GPU if available, else CPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Optional: keep track of the type for AMP autocast
+device_type = 'cuda' if device == 'cuda' else 'cpu'
+
+# Now build the config
+vocab_size = sum(1 for _ in open("/content/dict.txt", encoding="utf-8"))
+config = GPTConfig(
+    vocab_size=vocab_size,
+    block_size=128,    # or 1024 for BioGPT-scale training
+    n_layer=6,         # change to 24 for BioGPT-size
+    n_head=6,          # change to 16 for BioGPT-size
+    n_embd=384,        # change to 1024 for BioGPT-size
+    dropout=0.1,
+    bias=True
+)
+
+# Create model and move to device
+model = GPT(config).to(device)
+print("Params (M):", sum(p.numel() for p in model.parameters())/1e6)
+
+print(vocab_size)
+
+"""### 1.9 Define loss function"""
+
+def estimate_loss(model):
+    out = {}
+    model.eval()
+    with torch.inference_mode():
+        for split in ['train', 'valid']:
+            losses = torch.zeros(eval_iters)
+            for k in range(eval_iters):
+                X, Y = get_batch(split)
+                with ctx:
+                    logits, loss = model(X, Y)
+                losses[k] = loss.item()
+            out[split] = losses.mean()
+    model.train()
+    return out
+
+"""### 1.10 Define the training configuration"""
+
+# Training Config
+import torch
+from contextlib import nullcontext
+
+learning_rate = 1e-4 #more stable training, earlier 1e-4
+max_iters = 120000 #increase from 25000
+warmup_steps = 1000 #smoother initial train, earlier 100
+min_lr = 5e-4 #lower rate, earlier 5e-4
+eval_iters = 500 # increased from 100
+batch_size = 32 # changed from 16, better gradient estimate
+block_size = 128 #changed from 64, capture longer range dependencies
+
+gradient_accumulation_steps = 32 # reduced from 50
+
+device =  "cuda" if torch.cuda.is_available() else "cpu"
+device_type = 'cuda' if 'cuda' in device else 'cpu' # for later use in torch.autocast
+# note: float16 data type will automatically use a GradScaler
+
+# How to use autocast https://wandb.ai/wandb_fc/tips/reports/How-To-Use-Autocast-in-PyTorch--VmlldzoyMTk4NTky
+#dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16' # 'float32', 'bfloat16', or 'float16', the latter will auto implement a GradScaler
+dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16' # 'float32', 'bfloat16', or 'float16', the latter will auto implement a GradScaler
+ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
+
+ctx = nullcontext() if device_type == 'cpu' else torch.amp.autocast(device_type=device_type, dtype=ptdtype)
+
+torch.set_default_device(device)
+torch.manual_seed(42)
+
+"""### 1.11 Define optimizers and learning rate"""
+
+from torch.optim.lr_scheduler import LinearLR,SequentialLR, CosineAnnealingLR
+
+##PUT IN WEIGHT DECAY, CHANGED BETA2 to 0.95
+optimizer =  torch.optim.AdamW(model.parameters(), lr=learning_rate, betas=(0.9, 0.95), weight_decay=0.1, eps=1e-9) #weight decay for regularization
+
+scheduler_warmup = LinearLR(optimizer, total_iters = warmup_steps) #Implement linear warmup
+scheduler_decay = CosineAnnealingLR(optimizer,T_max = max_iters - warmup_steps, eta_min = min_lr) #Implement lr decay
+scheduler = SequentialLR(optimizer, schedulers=[scheduler_warmup, scheduler_decay], milestones=[warmup_steps]) #Switching from warmup to decay
+
+# https://stackoverflow.com/questions/72534859/is-gradscaler-necessary-with-mixed-precision-training-with-pytorch
+scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))
+
+"""### 1.12 Run pre-training!"""
+
+best_val_loss = float('inf')
+best_model_params_path = "best_model_params.pt"
+train_loss_list, validation_loss_list = [], []
+
+# Ensure model is on the correct device
+model = model.to(device)
+
+# In your training loop
+for epoch in tqdm(range(max_iters)):
+    if epoch % eval_iters == 0 and epoch != 0:
+        # Ensure estimate_loss uses the correct device
+        losses = estimate_loss(model)
+        print(f"Epoch {epoch}: train loss {losses['train']:.4f}, val loss {losses['valid']:.4f}")
+        print(f"The current learning rate: {optimizer.param_groups[0]['lr']:.5f}")
+        train_loss_list += [losses['train']]
+        validation_loss_list += [losses['valid']]
+
+        if losses['valid'] < best_val_loss:
+            best_val_loss = losses['valid']
+            torch.save(model.state_dict(), best_model_params_path)
+
+    # Ensure X and y are on the correct device
+    X, y = get_batch("train")
+    X, y = X.to(device), y.to(device)
+
+    with ctx:
+        logits, loss = model(X, y)
+        loss = loss / gradient_accumulation_steps
+        scaler.scale(loss).backward()
+
+    if ((epoch + 1) % gradient_accumulation_steps == 0) or (epoch + 1 == max_iters):
+        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.5)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
+    scheduler.step()
+
+"""### 1.13 Plot training and validation losses"""
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+eval_every = eval_iters  # e.g., 500
+
+# Convert each tensor to float on CPU
+train_loss_np = [float(t.cpu()) for t in train_loss_list]
+valid_loss_np = [float(t.cpu()) for t in validation_loss_list]
+
+steps = np.arange(1, len(train_loss_np) + 1) * eval_every
+
+plt.figure(figsize=(6,4))
+plt.plot(steps, train_loss_np, label='train')
+plt.plot(steps, valid_loss_np, label='valid')
+plt.xlabel('Iteration')
+plt.ylabel('Loss')
+plt.title('Pretraining loss')
+plt.legend()
+plt.grid(True, alpha=0.3)
+plt.show()
+
+import torch
+
+ckpt_path = "best_model_params.pt"  # you saved this in the loop
+model.load_state_dict(torch.load(ckpt_path, map_location=device))
+model.eval()
+
+"""### 1.14 Evaluation on HoC Part 1 (the Hallmarks of Cancers corpus) classification dataset"""
+
+import os
+import pandas as pd
+from datasets import load_dataset
+from tqdm.auto import tqdm
+
+def download_and_save_hoc_splits(target_dir="/content/hoc"):
+    """
+    Downloads the bigbio/hallmarks_of_cancer dataset from Hugging Face,
+    formats it, and saves it as train.tsv, valid.tsv, and test.tsv
+    in the specified directory.
+
+    Args:
+        target_dir (str): The directory to save the .tsv files.
+    """
+    print("Downloading bigbio/hallmarks_of_cancer dataset...")
+    try:
+        # Load the dataset splits
+        train_data = load_dataset("bigbio/hallmarks_of_cancer", split="train")
+        valid_data = load_dataset("bigbio/hallmarks_of_cancer", split="validation")
+        test_data = load_dataset("bigbio/hallmarks_of_cancer", split="test")
+        print("Dataset downloaded successfully.")
+    except Exception as e:
+        print(f"Error downloading dataset: {e}")
+        print("Please ensure you have internet access and the 'datasets' library is installed (`pip install datasets`).")
+        return
+
+    os.makedirs(target_dir, exist_ok=True)
+    print(f"Ensured target directory exists: {target_dir}")
+
+    splits = {
+        "train": train_data,
+        "valid": valid_data,
+        "test": test_data,
+    }
+
+    for split_name, dataset in splits.items():
+        output_path = os.path.join(target_dir, f"{split_name}.tsv")
+        print(f"Processing '{split_name}' split and saving to {output_path}...")
+
+        processed_data = []
+        # Iterate with tqdm for progress bar
+        for item in tqdm(dataset, desc=f"Processing {split_name}", leave=False):
+            text = item.get("text", "")
+            labels_list = item.get("labels", [])
+
+            # Handle the [' none '] case and join the list into a string
+            # Using '; ' as a separator, similar to how multi-label strings might appear
+            if labels_list == [' none '] or not labels_list:
+                label_str = "" # Represent 'none' or empty list as an empty string
+            else:
+                # Filter out ' none ' if mixed with others, though unlikely based on dataset viewer
+                valid_labels = [lbl for lbl in labels_list if lbl.strip().lower() != 'none']
+                label_str = "; ".join(valid_labels) # Join valid labels with a separator
+
+            # Append as a dictionary for easy DataFrame creation later
+            # Replace tabs and newlines in text to avoid breaking TSV format
+            cleaned_text = " ".join(text.split())
+            processed_data.append({"text": cleaned_text, "label": label_str})
+
+        # Convert to DataFrame and save as TSV
+        if processed_data:
+            df = pd.DataFrame(processed_data)
+            # Ensure columns are in the order expected by load_hoc_tsv heuristic (text, label)
+            df = df[["text", "label"]]
+            df.to_csv(output_path, sep="\t", index=False, header=False) # Save without index and header
+            print(f"Successfully saved {output_path}")
+        else:
+            print(f"No data processed for split '{split_name}'.")
+
+    print("\nDataset processing complete.")
+
+# Commented out IPython magic to ensure Python compatibility.
+# ===== Zero-shot HoC evaluation for your PRE-TRAINED GPT (with cue + EOS delay) =====
+# Uses your existing GPT / GPTConfig and loads ckpt_path="best_model_params.pt"
+
+# installs
+!pip -q install sacremoses==0.0.53 scikit-learn==1.5.1
+
+import os, math, difflib, tempfile, subprocess
+import numpy as np
+import pandas as pd
+from tqdm.auto import tqdm
+import torch
+import torch.nn.functional as F
+from sklearn.metrics import precision_recall_fscore_support
+from sacremoses import MosesDetokenizer
+
+# ---------- paths ----------
+HOC_DIR     = "/content/hoc"
+download_and_save_hoc_splits(HOC_DIR)             # train.tsv / valid.tsv / test.tsv live here
+BPE_CODES   = "/content/bpecodes"            # from BioGPT
+DICT_TXT    = "/content/dict.txt"            # from BioGPT
+FASTBPE_BIN = "/content/fastBPE/fast"        # compiled earlier
+ckpt_path   = ckpt_path if 'ckpt_path' in globals() else "best_model_params.pt"
+
+os.makedirs(HOC_DIR, exist_ok=True)
+
+# ---------- ensure fastBPE + BioGPT codes/dict ----------
+if not os.path.exists(FASTBPE_BIN):
+    !git clone -q https://github.com/glample/fastBPE.git /content/fastBPE
+#     %cd /content/fastBPE
+    !g++ -std=c++11 -O3 -pthread fastBPE/main.cc -IfastBPE -o fast
+#     %cd /content
+if not os.path.exists(BPE_CODES):
+    !wget -q -O /content/bpecodes https://raw.githubusercontent.com/microsoft/BioGPT/main/data/BioGPT/bpecodes
+if not os.path.exists(DICT_TXT):
+    !wget -q -O /content/dict.txt  https://raw.githubusercontent.com/microsoft/BioGPT/main/data/BioGPT/dict.txt
+
+# ---------- vocab maps ----------
+token2id, id2token = {}, {}
+with open(DICT_TXT, encoding="utf-8") as f:
+    for i, line in enumerate(f):
+        tok = line.split()[0]
+        token2id[tok] = i
+        id2token[i] = tok
+
+eos_id = token2id.get("</s>", 0)
+pad_id = eos_id  # safe pad; loss is masked anyway
+
+# ---------- BPE helpers ----------
+def bpe_encode_lines(lines, shard_size=2000, desc="BPE"):
+    if len(lines) == 0:
+        return []
+    out = []
+    with tempfile.TemporaryDirectory() as td:
+        for start in tqdm(range(0, len(lines), shard_size), desc=f"{desc} ({len(lines)} lines)", leave=False):
+            chunk = lines[start:start+shard_size]
+            src = os.path.join(td, f"src_{start}.txt")
+            dst = os.path.join(td, f"dst_{start}.bpe")
+            with open(src, "w", encoding="utf-8") as w:
+                for s in chunk: w.write((s or "").strip() + "\n")
+            subprocess.check_call([FASTBPE_BIN, "applybpe", dst, src, BPE_CODES])
+            with open(dst, "r", encoding="utf-8") as r:
+                for line in r:
+                    out.append(line.strip().split())
+    return out
+
+def tokens_to_ids(bpe_tokens):
+    ids = []
+    for t in bpe_tokens:
+        ids.append(token2id.get(t, pad_id))
+    return ids, 0
+
+def bpe_decode_tokens(bpe_tokens):
+    s = ' '.join(bpe_tokens).replace('@@ ', '')
+    return MosesDetokenizer(lang='en').detokenize(s.split())
+
+# ---------- load HoC test ----------
+def load_hoc_tsv(path):
+    df = pd.read_csv(path, sep="\t", header=None, dtype=str).fillna("")
+    assert df.shape[1] == 2, f"{path} must have 2 columns"
+    avg0, avg1 = df[0].astype(str).str.len().mean(), df[1].astype(str).str.len().mean()
+    df.columns = ["text","label"] if avg0 > avg1 else ["label","text"]
+    return df
+
+test_path = os.path.join(HOC_DIR, "test.tsv")
+assert os.path.exists(test_path), f"Missing {test_path}"
+test_df = load_hoc_tsv(test_path)
+print("Test size:", len(test_df))
+
+# ---------- the 10 Hallmarks (no 'empty') ----------
+HALLMARKS = [
+    "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",
+]
+
+def split_labels(s: str):
+    s = (s or "").strip()
+    if not s: return []
+    for sep in [",",";","|"]:
+        if sep in s:
+            return [p.strip() for p in s.split(sep) if p.strip()]
+    return [s]
+
+def normalize_labels(labs):
+    keep, low = [], [L.lower() for L in HALLMARKS]
+    for x in labs:
+        xl = x.lower().strip()
+        if xl in low:
+            keep.append(HALLMARKS[low.index(xl)])
+        else:
+            best = difflib.get_close_matches(xl, low, n=1, cutoff=0.7)
+            if best:
+                keep.append(HALLMARKS[low.index(best[0])])
+    return sorted(dict.fromkeys(keep))
+
+# ---------- Build allowed-token mask (labels + separators + </s>) & first-step forbids ----------
+def build_allowed_mask_and_first_forbid(vocab_size, device):
+    allowed = set()
+    sep_ids = set()
+    # Hallmark tokens (all tokens that appear in these strings)
+    for bpe in bpe_encode_lines(HALLMARKS, desc="BPE hallmarks"):
+        ids, _ = tokens_to_ids(bpe); allowed.update(ids)
+    # Separators; we also record their token ids to block at the first step
+    SEPS = [", ", ",", "; ", ";", "|", " and "]
+    for sep in SEPS:
+        bpe = bpe_encode_lines([sep], desc="BPE seps")[0]
+        ids, _ = tokens_to_ids(bpe)
+        allowed.update(ids)
+        sep_ids.update(ids)
+    allowed.add(eos_id)
+
+    mask = torch.full((vocab_size,), float('-inf'), device=device)
+    mask[list(allowed)] = 0.0
+    first_forbid = torch.zeros((vocab_size,), dtype=torch.bool, device=device)
+    first_forbid[list(sep_ids)] = True
+    first_forbid[eos_id] = True  # never allow EOS as the first generated token
+    return mask, first_forbid
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+ALLOWED_MASK, FIRST_STEP_FORBID = build_allowed_mask_and_first_forbid(len(token2id), device)
+
+# ---------- Build contexts (text </s> + textual cue) ----------
+PROMPT_TEXT = " hallmarks of cancer:"  # small cue after abstract
+PROMPT_BPE  = bpe_encode_lines([PROMPT_TEXT], desc="BPE prompt")[0]
+PROMPT_IDS, _ = tokens_to_ids(PROMPT_BPE)
+
+def make_context_with_prompt(df):
+    texts = df["text"].astype(str).tolist()
+    bpes  = bpe_encode_lines(texts, desc="BPE test ctx")
+    ctx = []
+    for bpe in bpes:
+        ids, _ = tokens_to_ids(bpe)
+        ctx.append(np.array(ids + [eos_id] + PROMPT_IDS, dtype=np.int64))
+    return ctx
+
+def pad_batch(seqs):
+    L = max(len(s) for s in seqs)
+    out = np.full((len(seqs), L), pad_id, dtype=np.int64)
+    for i, s in enumerate(seqs):
+        out[i, :len(s)] = s
+    return torch.from_numpy(out)
+
+def ids_to_tokens(ids):
+    return [id2token.get(int(i), "<unk>") for i in ids]
+
+def to_canonical(pred_chunk: str):
+    s = (pred_chunk or "").strip().lower()
+    low = [L.lower() for L in HALLMARKS]
+    if s in low: return HALLMARKS[low.index(s)]
+    best = difflib.get_close_matches(s, low, n=1, cutoff=0.7)
+    return HALLMARKS[low.index(best[0])] if best else None
+
+# ---------- Require your GPT & GPTConfig from pretraining ----------
+assert 'GPT' in globals(), "Please define your GPT class (same as pretraining) before running this cell."
+assert 'GPTConfig' in globals(), "Please ensure GPTConfig is defined."
+
+cfg = GPTConfig(
+    vocab_size=len(token2id),
+    block_size=(config.block_size if 'config' in globals() else 128),
+    n_layer=(config.n_layer if 'config' in globals() else 6),
+    n_head=(config.n_head if 'config' in globals() else 6),
+    n_embd=(config.n_embd if 'config' in globals() else 384),
+    dropout=(config.dropout if 'config' in globals() else 0.1),
+    bias=(config.bias if 'config' in globals() else True),
+)
+base = GPT(cfg).to(device)
+
+# safe WPE resize when loading the checkpoint
+def load_with_wpe_resize(model, ckpt_path):
+    sd = torch.load(ckpt_path, map_location="cpu")
+    key = "transformer.wpe.weight"
+    if key in sd:
+        old = sd[key]
+        new_w = model.transformer.wpe.weight
+        new_len = new_w.shape[0]
+        if old.shape[0] != new_len:
+            new = new_w.data.clone()
+            n = min(new_len, old.shape[0])
+            new[:n] = old[:n]
+            if new_len > n:
+                torch.nn.init.normal_(new[n:], mean=0.0, std=0.02)
+            sd[key] = new
+    missing, unexpected = base.load_state_dict(sd, strict=False)
+    if missing or unexpected:
+        print("Missing keys:", missing)
+    print("Loaded PRETRAINED checkpoint:", ckpt_path)
+
+assert os.path.exists(ckpt_path), f"Checkpoint not found: {ckpt_path}"
+load_with_wpe_resize(base, ckpt_path)
+base.eval()
+
+# ---------- Constrained greedy decode with cue + EOS delay ----------
+@torch.no_grad()
+def gpt_generate_with_cue(model, idx, allowed_mask, first_step_forbid,
+                          max_new_tokens=24, min_new_before_eos=2, eos_penalty=-2.0, temperature=0.0):
+    """
+    - Restrict vocabulary with `allowed_mask`
+    - For the very first generated token, forbid separators + EOS
+    - For the first `min_new_before_eos` tokens, disallow EOS entirely
+    - After that, add a small penalty to EOS (so it doesn't end too early)
+    """
+    out = idx.clone()
+    B = out.size(0)
+    finished = torch.zeros(B, dtype=torch.bool, device=out.device)
+    steps = 0
+    for _ in range(max_new_tokens):
+        ctx = out[:, -model.config.block_size:]
+        logits, _ = model(ctx)                 # (B,1,V)
+        logits = logits[:, -1, :]              # (B,V)
+
+        # restrict to label vocab
+        logits = logits + allowed_mask
+
+        # first token: block separators + EOS
+        if steps == 0:
+            logits[:, first_step_forbid] = -1e9
+
+        # delay EOS for a couple steps, then mildly penalize
+        if steps < min_new_before_eos:
+            logits[:, eos_id] = -1e9
+        else:
+            logits[:, eos_id] += eos_penalty
+
+        # pick next
+        if temperature <= 0:
+            next_id = torch.argmax(logits, dim=-1)
+        else:
+            probs = F.softmax(logits / temperature, dim=-1)
+            next_id = torch.multinomial(probs, num_samples=1).squeeze(1)
+
+        next_id = next_id.masked_fill(finished, eos_id)
+        out = torch.cat([out, next_id.unsqueeze(1)], dim=1)
+        finished |= (next_id == eos_id)
+        steps += 1
+        if bool(finished.all()):
+            break
+    return out[:, idx.size(1):]
+
+@torch.no_grad()
+def predict_labels_for_batch_generative(xb):
+    gens = gpt_generate_with_cue(
+        base, xb, allowed_mask=ALLOWED_MASK, first_step_forbid=FIRST_STEP_FORBID,
+        max_new_tokens=24, min_new_before_eos=2, eos_penalty=-2.0, temperature=0.0
+    )
+    preds = []
+    for g in gens:
+        toks = ids_to_tokens(g.detach().cpu().numpy())
+        toks = toks[: toks.index("</s>")] if "</s>" in toks else toks
+        label_str = bpe_decode_tokens(toks).strip().lower()
+
+        parts = []
+        for sep in [",",";","|"]:
+            if sep in label_str:
+                parts = [p.strip() for p in label_str.split(sep) if p.strip()]
+                break
+        if not parts:
+            parts = [label_str] if label_str else []
+
+        mapped = []
+        for p in parts:
+            can = to_canonical(p)
+            if can and can not in mapped:
+                mapped.append(can)
+        preds.append(mapped)  # may be []
+    return preds
+
+# ---------- Run decoding on TEST ----------
+ctx_test = make_context_with_prompt(test_df)
+preds_all = []
+B = 32
+for i in tqdm(range(0, len(ctx_test), B), desc="Decoding (pretrain+cue, test)"):
+    xb = pad_batch(ctx_test[i:i+B]).to(device)
+    preds_all.extend(predict_labels_for_batch_generative(xb))
+
+# ---------- Ground truth & metrics (10 hallmarks only) ----------
+y_true = [ normalize_labels(split_labels(s)) for s in test_df["label"].astype(str).tolist() ]
+LABELS = HALLMARKS
+LIDX = {l:i for i,l in enumerate(LABELS)}
+def binarize(labs):
+    v = [0]*len(LABELS)
+    for l in labs:
+        if l in LIDX: v[LIDX[l]] = 1
+    return v
+
+Y_true = np.array([binarize(l) for l in y_true], dtype=np.int64)
+Y_pred = np.array([binarize(l) for l in preds_all], dtype=np.int64)
+
+micro_p, micro_r, micro_f1, _ = precision_recall_fscore_support(Y_true, Y_pred, average='micro', zero_division=0)
+macro_p, macro_r, macro_f1, _ = precision_recall_fscore_support(Y_true, Y_pred, average='macro', zero_division=0)
+
+print(f"\n[PRETRAIN+cue] HALLMARKS-ONLY  Micro  P/R/F1: {micro_p:.4f} / {micro_r:.4f} / {micro_f1:.4f}")
+print(  f"[PRETRAIN+cue] HALLMARKS-ONLY  Macro  P/R/F1: {macro_p:.4f} / {macro_r:.4f} / {macro_f1:.4f}")
+
+perclass = precision_recall_fscore_support(Y_true, Y_pred, average=None, zero_division=0)
+per_df_pre = pd.DataFrame({
+    "label": LABELS,
+    "precision": perclass[0],
+    "recall":    perclass[1],
+    "f1":        perclass[2],
+    "support":   perclass[3],
+}).sort_values("label")
+
+print("\nPer-class results (PRETRAIN+cue, 10 hallmarks):")
+print(per_df_pre.to_string(index=False))
+
+per_df_pre.to_csv("hoc_test_results_pretrain_cue.csv", index=False)
+print("Saved: hoc_test_results_pretrain_cue.csv")
+
+# (optional) exclude empty-label rows from eval:
+# mask = (Y_true.sum(axis=1) > 0)
+# ... recompute scores on Y_true[mask], Y_pred[mask]
+
+"""### 1.15 Evaluation on HoC Part 2 (the Hallmarks of Cancers corpus) classification dataset"""
+
+# === Show 10 "questions" (abstract + prompt) and the model's answers (pretrained+cue) ===
+import os, difflib, numpy as np, pandas as pd, torch, torch.nn.functional as F
+from tqdm.auto import tqdm
+from sklearn.metrics import precision_recall_fscore_support
+
+# ---- Assumptions / fallbacks ----
+HOC_DIR   = globals().get("HOC_DIR", "/content/hoc")
+ckpt_path = globals().get("ckpt_path", "best_model_params.pt")
+device    = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Hallmarks (10 classes, no "empty")
+HALLMARKS = [
+    "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",
+]
+
+# ---------- Helper fallbacks if not defined earlier ----------
+def _need(name): return name not in globals()
+
+# TSV loader
+if _need("load_hoc_tsv"):
+    def load_hoc_tsv(path):
+        df = pd.read_csv(path, sep="\t", header=None, dtype=str).fillna("")
+        assert df.shape[1] == 2, f"{path} must have 2 columns"
+        avg0, avg1 = df[0].astype(str).str.len().mean(), df[1].astype(str).str.len().mean()
+        df.columns = ["text","label"] if avg0 > avg1 else ["label","text"]
+        return df
+
+# If test_df not in memory, load it
+if "test_df" not in globals():
+    test_df = load_hoc_tsv(os.path.join(HOC_DIR, "test.tsv"))
+
+# Simple label split/normalization utilities
+def split_labels(s: str):
+    s = (s or "").strip()
+    if not s: return []
+    for sep in [",",";","|"]:
+        if sep in s:
+            return [p.strip() for p in s.split(sep) if p.strip()]
+    return [s]
+
+def normalize_labels(labs):
+    keep, low = [], [L.lower() for L in HALLMARKS]
+    for x in labs:
+        xl = x.lower().strip()
+        if xl in low:
+            keep.append(HALLMARKS[low.index(xl)])
+        else:
+            best = difflib.get_close_matches(xl, low, n=1, cutoff=0.7)
+            if best:
+                keep.append(HALLMARKS[low.index(best[0])])
+    # de-dup & stable order
+    seen, out = set(), []
+    for k in keep:
+        if k not in seen:
+            seen.add(k); out.append(k)
+    return out
+
+# BPE helpers (must exist: token2id, id2token, bpe_encode_lines, tokens_to_ids, bpe_decode_tokens, eos_id, pad_id)
+for req in ["token2id","id2token","bpe_encode_lines","tokens_to_ids","bpe_decode_tokens","eos_id","pad_id"]:
+    assert req in globals(), f"Missing `{req}` — run the setup cell that defines dict/bpecodes and BPE helpers."
+
+# Build allowed-token mask & first-step forbids if not present
+if _need("ALLOWED_MASK") or _need("FIRST_STEP_FORBID"):
+    def build_allowed_mask_and_first_forbid(vocab_size, device):
+        allowed = set(); sep_ids = set()
+        # all tokens that appear in hallmark strings
+        for bpe in bpe_encode_lines(HALLMARKS, desc="BPE hallmarks"):
+            ids, _ = tokens_to_ids(bpe); allowed.update(ids)
+        # separators (also block them on very first generated step)
+        SEPS = [", ", ",", "; ", ";", "|", " and "]
+        for sep in SEPS:
+            bpe = bpe_encode_lines([sep], desc="BPE seps")[0]
+            ids, _ = tokens_to_ids(bpe); allowed.update(ids); sep_ids.update(ids)
+        allowed.add(eos_id)
+        mask = torch.full((vocab_size,), float('-inf'), device=device)
+        mask[list(allowed)] = 0.0
+        first_forbid = torch.zeros((vocab_size,), dtype=torch.bool, device=device)
+        first_forbid[list(sep_ids)] = True
+        first_forbid[eos_id] = True
+        return mask, first_forbid
+    ALLOWED_MASK, FIRST_STEP_FORBID = build_allowed_mask_and_first_forbid(len(token2id), device)
+
+# Prompt (the "question" cue)
+PROMPT_TEXT = " hallmarks of cancer:"
+PROMPT_BPE  = bpe_encode_lines([PROMPT_TEXT], desc="BPE prompt")[0]
+PROMPT_IDS, _ = tokens_to_ids(PROMPT_BPE)
+
+# Build contexts with prompt
+def make_context_with_prompt(rows):
+    bpes = bpe_encode_lines(rows["text"].astype(str).tolist(), desc="BPE ctx (sample)")
+    ctx = []
+    for bpe in bpes:
+        ids, _ = tokens_to_ids(bpe)
+        ctx.append(np.array(ids + [eos_id] + PROMPT_IDS, dtype=np.int64))
+    return ctx
+
+def pad_batch(seqs):
+    L = max(len(s) for s in seqs)
+    out = np.full((len(seqs), L), pad_id, dtype=np.int64)
+    for i, s in enumerate(seqs):
+        out[i, :len(s)] = s
+    return torch.from_numpy(out)
+
+def ids_to_tokens(ids):
+    return [id2token.get(int(i), "<unk>") for i in ids]
+
+def to_canonical(pred_chunk: str):
+    s = (pred_chunk or "").strip().lower()
+    low = [L.lower() for L in HALLMARKS]
+    if s in low: return HALLMARKS[low.index(s)]
+    best = difflib.get_close_matches(s, low, n=1, cutoff=0.7)
+    return HALLMARKS[low.index(best[0])] if best else None
+
+# If the pretrained model (`base`) isn’t loaded yet, load it
+if _need("base"):
+    assert 'GPT' in globals() and 'GPTConfig' in globals(), "Define GPT and GPTConfig first (your pretraining classes)."
+    assert os.path.exists(ckpt_path), f"Checkpoint not found: {ckpt_path}"
+    cfg = GPTConfig(
+        vocab_size=len(token2id),
+        block_size=(config.block_size if 'config' in globals() else 128),
+        n_layer=(config.n_layer if 'config' in globals() else 6),
+        n_head=(config.n_head if 'config' in globals() else 6),
+        n_embd=(config.n_embd if 'config' in globals() else 384),
+        dropout=(config.dropout if 'config' in globals() else 0.1),
+        bias=(config.bias if 'config' in globals() else True),
+    )
+    base = GPT(cfg).to(device)
+    # safe WPE resize
+    def load_with_wpe_resize(model, path):
+        sd = torch.load(path, map_location="cpu")
+        key = "transformer.wpe.weight"
+        if key in sd:
+            old = sd[key]
+            new_w = model.transformer.wpe.weight
+            new_len = new_w.shape[0]
+            if old.shape[0] != new_len:
+                new = new_w.data.clone()
+                n = min(new_len, old.shape[0])
+                new[:n] = old[:n]
+                if new_len > n:
+                    torch.nn.init.normal_(new[n:], mean=0.0, std=0.02)
+                sd[key] = new
+        model.load_state_dict(sd, strict=False)
+    load_with_wpe_resize(base, ckpt_path)
+    base.eval()
+
+# Constrained generation with cue + EOS delay (define if missing)
+if _need("gpt_generate_with_cue"):
+    @torch.no_grad()
+    def gpt_generate_with_cue(model, idx, allowed_mask, first_step_forbid,
+                              max_new_tokens=24, min_new_before_eos=2, eos_penalty=-2.0, temperature=0.0):
+        out = idx.clone()
+        B = out.size(0)
+        finished = torch.zeros(B, dtype=torch.bool, device=out.device)
+        steps = 0
+        for _ in range(max_new_tokens):
+            ctx = out[:, -model.config.block_size:]
+            logits, _ = model(ctx)           # (B,1,V)
+            logits = logits[:, -1, :]        # (B,V)
+            logits = logits + allowed_mask    # restrict vocab
+            if steps == 0:
+                logits[:, first_step_forbid] = -1e9
+            if steps < min_new_before_eos:
+                logits[:, eos_id] = -1e9
+            else:
+                logits[:, eos_id] += eos_penalty
+            if temperature <= 0:
+                next_id = torch.argmax(logits, dim=-1)
+            else:
+                probs = F.softmax(logits / temperature, dim=-1)
+                next_id = torch.multinomial(probs, num_samples=1).squeeze(1)
+            next_id = next_id.masked_fill(finished, eos_id)
+            out = torch.cat([out, next_id.unsqueeze(1)], dim=1)
+            finished |= (next_id == eos_id)
+            steps += 1
+            if bool(finished.all()):
+                break
+        return out[:, idx.size(1):]
+
+# ---------- Sample 10 and print Q&A ----------
+SAMPLE_N = 10
+sample = test_df.sample(n=min(SAMPLE_N, len(test_df)), random_state=42).reset_index(drop=True)
+
+# prepare contexts
+ctx = make_context_with_prompt(sample)
+B = 10  # single batch is fine here
+xb = pad_batch(ctx).to(device)
+
+# generate
+gens = gpt_generate_with_cue(
+    base, xb, allowed_mask=ALLOWED_MASK, first_step_forbid=FIRST_STEP_FORBID,
+    max_new_tokens=24, min_new_before_eos=2, eos_penalty=-2.0, temperature=0.0
+)
+
+# decode + print
+for i, g in enumerate(gens):
+    text = sample.loc[i, "text"]
+    gold = normalize_labels(split_labels(sample.loc[i, "label"]))
+
+    toks = ids_to_tokens(g.detach().cpu().numpy())
+    toks = toks[: toks.index("</s>")] if "</s>" in toks else toks
+    raw = ' '.join(toks).replace('@@ ', '').strip().lower()
+
+    # split raw into parts and map to canonical labels
+    parts = []
+    for sep in [",",";","|"]:
+        if sep in raw:
+            parts = [p.strip() for p in raw.split(sep) if p.strip()]
+            break
+    if not parts:
+        parts = [raw] if raw else []
+    pred = []
+    for p in parts:
+        can = to_canonical(p)
+        if can and can not in pred:
+            pred.append(can)
+
+    print(f"\n=== Example {i+1} ===")
+    print("QUESTION:")
+    print("Abstract:", (text.replace("\n"," ")[:350] + ("..." if len(text) > 350 else "")))
+    print("Prompt:  hallmarks of cancer:")
+    print("GOLD:   ", gold if gold else "[]")
+    print("ANSWER: ", pred if pred else "[]")
+    print("Raw gen:", raw if raw else "<empty>")
+
+"""## Part 2: Finetuning
+
+### 2.1 Setup: paths + installs
+"""
+
+# Commented out IPython magic to ensure Python compatibility.
+# --- Setup: paths + installs (run once) ---
+!pip -q install sacremoses==0.0.53 scikit-learn==1.5.1
+
+import os, subprocess, json, math, random, difflib, tempfile, shutil
+from pathlib import Path
+import numpy as np
+import pandas as pd
+from collections import Counter, defaultdict
+
+import torch, torch.nn as nn, torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+from torch.optim.lr_scheduler import LinearLR, SequentialLR, CosineAnnealingLR
+from sacremoses import MosesDetokenizer
+from tqdm.auto import tqdm  # <-- used in BPE w/ progress
+
+# ---- paths ----
+HOC_DIR   = "/content/hoc"                # << put your train/valid/test.tsv here
+BPE_CODES = "/content/bpecodes"           # from your pre-training cell
+DICT_TXT  = "/content/dict.txt"           # from your pre-training cell
+FASTBPE   = "/content/fastBPE/fast"       # compiled earlier in your notebook
+
+os.makedirs(HOC_DIR, exist_ok=True)
+
+# Ensure fastBPE exists (rebuild if needed)
+if not os.path.exists(FASTBPE):
+    !git clone -q https://github.com/glample/fastBPE.git /content/fastBPE
+#     %cd /content/fastBPE
+    !g++ -std=c++11 -O3 -pthread fastBPE/main.cc -IfastBPE -o fast
+#     %cd /content
+
+# ---- load BioGPT dictionary ----
+token2id = {}
+id2token = {}
+with open(DICT_TXT, encoding="utf-8") as f:
+    for i, line in enumerate(f):
+        tok = line.split()[0]
+        token2id[tok] = i
+        id2token[i] = tok
+
+# pick special ids
+eos_id = token2id.get("</s>", 0)
+pad_id = eos_id  # safe padding with eos for inputs; we mask loss anyway
+
+# ---- BPE encode/decode helpers (fastBPE uses '@@' continuation) ----
+def bpe_encode_lines(lines, shard_size=2000, desc="BPE"):
+    """
+    Progress-enabled BPE encoding using fastBPE, processing in shards.
+    Returns: list[list[str]] (BPE tokens per line)
+    """
+    if len(lines) == 0:
+        return []
+    out_tokens = []
+    with tempfile.TemporaryDirectory() as td:
+        for start in tqdm(range(0, len(lines), shard_size), desc=f"{desc} ({len(lines)} lines)", leave=False):
+            chunk = lines[start:start+shard_size]
+            src = os.path.join(td, f"src_{start}.txt")
+            dst = os.path.join(td, f"dst_{start}.bpe")
+            with open(src, "w", encoding="utf-8") as f:
+                for s in chunk:
+                    f.write((s or "").strip() + "\n")
+            subprocess.check_call([FASTBPE, "applybpe", dst, src, BPE_CODES])
+            with open(dst, "r", encoding="utf-8") as f:
+                for line in f:
+                    out_tokens.append(line.strip().split())
+    return out_tokens
+
+def bpe_decode_tokens(bpe_tokens):
+    """Merge '@@' continuations and detokenize to plain text (for label decoding)."""
+    s = ' '.join(bpe_tokens).replace('@@ ', '')
+    md = MosesDetokenizer(lang='en')
+    return md.detokenize(s.split())
+
+def tokens_to_ids(bpe_tokens):
+    ids = []
+    oov = 0
+    for t in bpe_tokens:
+        if t in token2id:
+            ids.append(token2id[t])
+        else:
+            ids.append(pad_id)  # unlikely, but safe fallback
+            oov += 1
+    return ids, oov
+
+"""### 2.2 Load HoC dataset and map targets to labels"""
+
+# --- Load HoC TSVs (2 columns, no header). Heuristically figure out which is text vs label. ---
+def load_hoc_tsv(path):
+    df = pd.read_csv(path, sep="\t", header=None, dtype=str).fillna("")
+    assert df.shape[1] == 2, f"Expected 2 columns in {path}, got {df.shape}"
+    avg0, avg1 = df[0].astype(str).str.len().mean(), df[1].astype(str).str.len().mean()
+    if avg0 > avg1:
+        df.columns = ["text", "label"]
+    else:
+        df.columns = ["label", "text"]
+    return df
+
+train_df = load_hoc_tsv(f"{HOC_DIR}/train.tsv")
+valid_df = load_hoc_tsv(f"{HOC_DIR}/valid.tsv")
+test_df  = load_hoc_tsv(f"{HOC_DIR}/test.tsv")
+
+print("Splits:", len(train_df), len(valid_df), len(test_df))
+
+# --- Hallmarks (10 classes; we ignore 'empty' for training and for reporting) ---
+HALLMARKS = [
+    "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",
+]
+
+def split_labels(s: str):
+    s = (s or "").strip()
+    if not s: return []
+    for sep in [",", ";", "|"]:
+        if sep in s:
+            return [p.strip() for p in s.split(sep) if p.strip()]
+    return [s]
+
+def normalize_labels(labs):
+    """Map raw labels (including fuzzy matches) to the 10 hallmarks; drop 'empty'."""
+    keep = []
+    low = [L.lower() for L in HALLMARKS]
+    for x in labs:
+        x_low = x.lower().strip()
+        if x_low in low:
+            keep.append(HALLMARKS[low.index(x_low)])
+        else:
+            best = difflib.get_close_matches(x_low, low, n=1, cutoff=0.7)
+            if best:
+                keep.append(HALLMARKS[low.index(best[0])])
+    # dedupe & sort for deterministic target text
+    return sorted(list(dict.fromkeys(keep)))
+
+def labels_to_target_text(labs):
+    labs = normalize_labels(labs)
+    if len(labs) == 0:
+        return None  # -> drop from training if empty-only
+    return ", ".join(labs)
+
+"""### 2.3 Redefine GPT architecture for full finetuning"""
+
+# --- Your GPT modules (same as in your pretraining code) ---
+class LayerNorm(nn.Module):
+    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):
+    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')
+        if not self.flash:
+            self.register_buffer("bias", torch.tril(torch.ones(config.block_size, config.block_size))
+                                       .view(1, 1, config.block_size, config.block_size))
+    def forward(self, x):
+        B, T, C = x.size()
+        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)
+        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2)
+        if self.flash:
+            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:
+            att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+            att = att.masked_fill(self.bias[:, :, :T, :T] == 0, float('-inf'))
+            att = F.softmax(att, dim=-1)
+            att = self.attn_dropout(att)
+            y = att @ v
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+        y = self.resid_dropout(self.c_proj(y))
+        return y
+
+class MLP(nn.Module):
+    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):
+        return self.dropout(self.c_proj(self.gelu(self.c_fc(x))))
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln1 = LayerNorm(config.n_embd, config.bias)
+        self.attn = CausalSelfAttention(config)
+        self.ln2 = LayerNorm(config.n_embd, 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
+
+from dataclasses import dataclass
+@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 GPT(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        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, config.bias),
+        ))
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        # weight tying
+        self.transformer.wte.weight = self.lm_head.weight
+
+        self.apply(self._init_weights)
+        for pn, p in self.named_parameters():
+            if pn.endswith('c_proj.weight'):
+                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):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None):
+        device = idx.device
+        B, T = idx.size()
+        assert T <= self.config.block_size
+        pos = torch.arange(0, T, dtype=torch.long, device=device)
+        tok_emb = self.transformer.wte(idx)
+        pos_emb = self.transformer.wpe(pos)
+        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:
+            logits = self.lm_head(x)               # (B,T,V)
+            loss = F.cross_entropy(
+                logits.view(-1, logits.size(-1)),
+                targets.view(-1),
+                ignore_index=-1
+            )
+            return logits, loss
+        else:
+            logits = self.lm_head(x[:, [-1], :])   # (B,1,V)
+            return logits, None
+
+"""### 2.4 Define Add SoftPrompt embeddings to input embeddings"""
+
+class GPTWithSoftPrompt(nn.Module):
+    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))
+        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
+
+        # 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 = self.soft_prompt.expand(B, -1, -1)                        # (B,P,C)
+        x = torch.cat([soft, x_tokens], dim=1)                           # (B,P+T,C)
+
+        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)
+
+        if targets is None:
+            # return next-token logits at last (standard for generation)
+            return logits[:, -1, :], None
+
+        # ----- FIX: next-token loss with soft-prompt padding -----
+        P = soft.size(1)
+        pad_ignore = torch.full((B, P), -1, dtype=targets.dtype, device=device)  # ignore for soft prompt
+        full_targets = torch.cat([pad_ignore, targets], dim=1)                    # (B,P+T)
+
+        # shift for next-token prediction
+        logits_lm  = logits[:, :-1, :].contiguous()                               # predict next token
+        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
+        )
+        return logits, loss
+
+"""### 2.5 Instantiate pre-training weights"""
+
+# --- Instantiate & (optionally) load your pretraining weights ---
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Use your pretrain block_size (128 in your earlier run). If different, the loader below can resize wpe.
+BLOCK_SIZE = 128  # set to 128 if that was your pretrain; otherwise set to your pretrain context length
+
+config = GPTConfig(
+    vocab_size=len(token2id),
+    block_size=BLOCK_SIZE,
+    n_layer=6, n_head=6, n_embd=384,
+    dropout=0.1, bias=True
+)
+base_gpt = GPT(config)
+
+def load_with_wpe_resize(model, ckpt_path):
+    sd = torch.load(ckpt_path, map_location="cpu")
+    key = "transformer.wpe.weight"
+    if key in sd:
+        old = sd[key]
+        new_len = model.transformer.wpe.weight.shape[0]
+        if old.shape[0] != new_len:
+            # copy existing, init the rest
+            new = model.transformer.wpe.weight.data.clone()
+            n = min(new_len, old.shape[0])
+            new[:n] = old[:n]
+            if new_len > n:
+                nn.init.normal_(new[n:], mean=0.0, std=0.02)
+            sd[key] = new
+    missing, unexpected = model.load_state_dict(sd, strict=False)
+    print("Loaded state dict with resize. Missing:", missing, "Unexpected:", unexpected)
+
+pt_path = "best_model_params.pt"
+if os.path.exists(pt_path):
+    load_with_wpe_resize(base_gpt, pt_path)
+    print("Loaded pretraining weights from:", pt_path)
+else:
+    print("No pretrain checkpoint found; training soft prompt from scratch on top of random GPT.")
+
+model = GPTWithSoftPrompt(base_gpt, prompt_len=1).to(device)
+
+"""### 2.6 Build a mask of token IDs that are allowed during generation"""
+
+# --- Constrained token mask (only hallmarks + separators + </s>) ---
+def build_allowed_token_mask(vocab_size, device):
+    allowed = set()
+    # hallmark token ids
+    for bpe in bpe_encode_lines(HALLMARKS, desc="BPE hallmarks"):
+        ids, _ = tokens_to_ids(bpe)
+        allowed.update(ids)
+    # separators
+    for sep in [", ", ",", "; ", ";", "|", " and "]:
+        bpe = bpe_encode_lines([sep], desc="BPE seps")[0]
+        ids, _ = tokens_to_ids(bpe)
+        allowed.update(ids)
+    allowed.add(eos_id)
+    mask = torch.full((vocab_size,), float('-inf'), device=device)
+    mask[list(allowed)] = 0.0
+    return mask
+
+ALLOWED_MASK = build_allowed_token_mask(len(token2id), device)
+
+"""### 2.7:
+
+- Define a dataset class that encodes abstracts and labels into token IDs (dropping empty-only rows for training if desired)
+- Concatenate them into input/target sequences respecting a block size
+- Provide a collate function to pad batches for training.
+"""
+
+# --- Dataset (drops empty-only rows for TRAIN to avoid collapse) ---
+class HoCGenDataset(Dataset):
+    def __init__(self, df, block_size=256, drop_empty_only=False, name=""):
+        self.block_size = block_size
+        self.samples = []
+
+        texts = df["text"].astype(str).tolist()
+        raw_labels = [split_labels(s) for s in df["label"].astype(str).tolist()]
+
+        # BPE encode texts with progress
+        text_bpe = bpe_encode_lines(texts, shard_size=2000, desc=f"BPE {name or 'dataset'}")
+
+        # Pre-encode unique label targets
+        targets = []
+        for labs in raw_labels:
+            tgt = labels_to_target_text(labs)  # None if empty-only
+            targets.append(tgt)
+        uniq_non_null = sorted(set([t for t in targets if t is not None]))
+
+        label_cache = {}
+        if len(uniq_non_null) > 0:
+            encoded = bpe_encode_lines(uniq_non_null, shard_size=200, desc=f"BPE labels {name or 'dataset'}")
+            for s, bpe in zip(uniq_non_null, encoded):
+                ids, _ = tokens_to_ids(bpe)
+                label_cache[s] = ids
+
+        # Pack samples
+        for bpe, tgt in tqdm(list(zip(text_bpe, targets)), total=len(text_bpe), desc=f"Packing {name or 'dataset'}", leave=False):
+            if drop_empty_only and tgt is None:
+                continue
+            text_ids, _ = tokens_to_ids(bpe)
+
+            if tgt is None:
+                label_ids = []
+            else:
+                label_ids = label_cache[tgt]
+
+            x_ids = text_ids + [eos_id]
+            y_ids = (label_ids + [eos_id]) if len(label_ids) > 0 else []
+
+            # respect block size
+            max_text = self.block_size - (2 if len(y_ids) > 0 else 1) - len(y_ids)
+            if max_text < 1:
+                x_ids = x_ids[:max(1, self.block_size // 2)]
+            else:
+                x_ids = x_ids[:max_text]
+
+            input_ids = x_ids + y_ids
+            targets_arr = ([-1] * len(x_ids)) + (y_ids if len(y_ids) > 0 else [])
+
+            self.samples.append((
+                np.array(input_ids, dtype=np.int64),
+                np.array(targets_arr, dtype=np.int64)
+            ))
+
+    def __len__(self): return len(self.samples)
+    def __getitem__(self, idx): return self.samples[idx]
+
+def collate(batch):
+    L = max(len(x[0]) for x in batch)
+    B = len(batch)
+    inputs  = np.full((B, L), pad_id, dtype=np.int64)
+    targets = np.full((B, L), -1,     dtype=np.int64)
+    for i, (inp, tgt) in enumerate(batch):
+        n = len(inp)
+        inputs[i, :n]  = inp
+        targets[i, :n] = tgt
+    return torch.from_numpy(inputs), torch.from_numpy(targets)
+
+"""### 2.8 Create dataloaders for the finetuning dataset"""
+
+# --- Datasets/Loaders ---
+BATCH_SIZE = 16
+
+# Train: drop empty-only rows (crucial)
+train_ds = HoCGenDataset(train_df, block_size=model.config.block_size, drop_empty_only=True,  name="train")
+# Valid: drop empty-only too (makes val loss meaningful)
+valid_ds = HoCGenDataset(valid_df, block_size=model.config.block_size, drop_empty_only=True,  name="valid")
+# Test: keep all rows; we'll evaluate on the 10 hallmarks only later
+test_ds  = HoCGenDataset(test_df,  block_size=model.config.block_size, drop_empty_only=False, name="test")
+
+cuda_gen = torch.Generator(device='cuda')  # or set a manual seed if you want
+
+train_loader = DataLoader(
+    train_ds, batch_size=BATCH_SIZE, shuffle=True,
+    collate_fn=collate, drop_last=True,
+    generator=cuda_gen,      # <-- key fix
+    pin_memory=True, pin_memory_device='cuda'
+)
+
+valid_loader = DataLoader(
+    valid_ds, batch_size=BATCH_SIZE, shuffle=False,
+    collate_fn=collate,
+    generator=cuda_gen,
+    pin_memory=True, pin_memory_device='cuda'
+)
+
+test_loader  = DataLoader(
+    test_ds,  batch_size=BATCH_SIZE, shuffle=False,
+    collate_fn=collate,
+    generator=cuda_gen,
+    pin_memory=True, pin_memory_device='cuda'
+)
+
+print(f"Train samples (non-empty only): {len(train_ds)}")
+print(f"Valid samples (non-empty only): {len(valid_ds)}")
+print(f"Test samples  (incl. empty):    {len(test_ds)}")
+
+xb, yb = next(iter(train_loader))
+assert (yb != -1).any(), "No supervised label tokens in this batch — are we dropping all rows?"
+
+xb, yb = xb.to(device), yb.to(device)
+with torch.no_grad():
+    _, loss = model(xb, yb)
+print("Initial loss:", float(loss))
+
+"""### 2.9
+
+- Feeds the current context into the model (self(ctx)).
+
+- Adds the allowed_mask to the logits so that only permitted token IDs (Hallmarks, separators, </s>) can be chosen; all others get -inf and are impossible to sample.
+
+- Picks the next token greedily (argmax) unless a temperature is set, in which case it samples.
+
+- Forces already finished sequences to emit </s> and stops early when all sequences are finished.
+"""
+
+# --- Constrained, batched decoding method for GPTWithSoftPrompt ---
+def constrained_generate_labels(self, idx, allowed_mask, max_new_tokens=24, temperature=0.0):
+    """
+    Batched decode. At each step, mask logits to the allowed set.
+    Returns only generated tail (B, Tgen).
+    """
+    self.eval()
+    B = idx.size(0)
+    out = idx.clone()
+    finished = torch.zeros(B, dtype=torch.bool, device=idx.device)
+
+    for _ in range(max_new_tokens):
+        ctx = out[:, -self.config.block_size:]
+        logits, _ = self(ctx)          # (B,V)
+        # apply constraint
+        logits = logits + allowed_mask
+        if temperature <= 0:
+            next_id = torch.argmax(logits, dim=-1)  # (B,)
+        else:
+            probs = F.softmax(logits / temperature, dim=-1)
+            next_id = torch.multinomial(probs, num_samples=1).squeeze(1)
+
+        next_id = next_id.masked_fill(finished, eos_id)
+        out = torch.cat([out, next_id.unsqueeze(1)], dim=1)
+        finished |= (next_id == eos_id)
+        if bool(finished.all()):
+            break
+    return out[:, idx.size(1):]
+
+# attach to instance/class
+GPTWithSoftPrompt.generate_labels = constrained_generate_labels
+
+"""### 2.10 Run the finetuning loop"""
+
+# --- Optimizer & schedulers (paper: 20k steps, warmup 1k, peak LR 1e-5) ---
+max_steps   = 20_000
+warmup      = 1_000
+peak_lr     = 1e-5
+eta_min     = 1e-6
+
+optimizer   = torch.optim.AdamW(model.parameters(), lr=peak_lr, betas=(0.9, 0.95), weight_decay=0.01, eps=1e-9)
+sched_warm  = LinearLR(optimizer, total_iters=warmup)
+sched_decay = CosineAnnealingLR(optimizer, T_max=max_steps - warmup, eta_min=eta_min)
+scheduler   = SequentialLR(optimizer, [sched_warm, sched_decay], milestones=[warmup])
+
+# AMP dtype: bf16 if supported, else fp16; enable GradScaler only if fp16
+amp_dtype = torch.bfloat16 if (torch.cuda.is_available() and torch.cuda.is_bf16_supported()) else torch.float16
+scaler    = torch.cuda.amp.GradScaler(enabled=(amp_dtype == torch.float16))
+
+def run_eval(loader):
+    model.eval()
+    losses = []
+    with torch.no_grad():
+        for xb, yb in tqdm(loader, desc="Valid", leave=False):
+            xb, yb = xb.to(device), yb.to(device)
+            with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=torch.cuda.is_available()):
+                _, loss = model(xb, yb)
+            losses.append(loss.item())
+    model.train()
+    return float(np.mean(losses)) if losses else 0.0
+
+# --- Training loop ---
+EVAL_EVERY = 500
+BEST_PATH  = "hoc_best.pt"
+
+best_val = float('inf')
+global_step = 0
+ema_loss = None
+pbar = tqdm(total=max_steps, desc="Finetuning (HoC)", leave=True)
+
+model.train()
+while global_step < max_steps:
+    for xb, yb in train_loader:
+        xb, yb = xb.to(device), yb.to(device)
+        with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=torch.cuda.is_available()):
+            _, loss = model(xb, yb)
+
+        scaler.scale(loss).backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
+        scaler.step(optimizer)
+        scaler.update()
+        optimizer.zero_grad(set_to_none=True)
+        scheduler.step()
+
+        global_step += 1
+        pbar.update(1)
+
+        cur = loss.item()
+        ema_loss = cur if ema_loss is None else (0.95 * ema_loss + 0.05 * cur)
+        pbar.set_postfix({
+            "train_loss": f"{cur:.3f}",
+            "ema":        f"{ema_loss:.3f}",
+            "best_val":   f"{best_val:.3f}" if best_val < float('inf') else "—",
+            "lr":         f"{optimizer.param_groups[0]['lr']:.2e}",
+        })
+
+        if global_step % EVAL_EVERY == 0:
+            val_loss = run_eval(valid_loader)
+            if val_loss < best_val:
+                best_val = val_loss
+                torch.save(model.state_dict(), BEST_PATH)
+            pbar.set_postfix({
+                "train_loss": f"{cur:.3f}",
+                "ema":        f"{ema_loss:.3f}",
+                "best_val":   f"{best_val:.3f}",
+                "lr":         f"{optimizer.param_groups[0]['lr']:.2e}",
+            })
+
+        if global_step >= max_steps:
+            break
+
+pbar.close()
+
+# reload best
+if os.path.exists(BEST_PATH):
+    model.load_state_dict(torch.load(BEST_PATH, map_location=device))
+    print("Loaded best checkpoint:", BEST_PATH, " (val_loss:", f"{best_val:.4f}", ")")
+
+"""### 2.11 Classification evaluation"""
+
+# --- Build context-only inputs (text </s>) directly from raw test_df ---
+def make_context_only(df):
+    texts = df["text"].astype(str).tolist()
+    bpes  = bpe_encode_lines(texts, desc="BPE test ctx")
+    ctx = []
+    for bpe in bpes:
+        ids, _ = tokens_to_ids(bpe)
+        ctx.append(np.array(ids + [eos_id], dtype=np.int64))
+    return ctx
+
+def pad_batch(seqs):
+    L = max(len(s) for s in seqs)
+    out = np.full((len(seqs), L), pad_id, dtype=np.int64)
+    for i, s in enumerate(seqs):
+        out[i, :len(s)] = s
+    return torch.from_numpy(out)
+
+def ids_to_tokens(ids):
+    return [id2token.get(int(i), "<unk>") for i in ids]
+
+def to_canonical(pred_chunk: str):
+    s = (pred_chunk or "").strip().lower()
+    low = [L.lower() for L in HALLMARKS]
+    if s in low:
+        return HALLMARKS[low.index(s)]
+    best = difflib.get_close_matches(s, low, n=1, cutoff=0.7)
+    return HALLMARKS[low.index(best[0])] if best else None
+
+def predict_labels_for_batch(xb):
+    """xb: (B, T) context-only input ids (text </s>)."""
+    with torch.no_grad():
+        gens = model.generate_labels(xb, allowed_mask=ALLOWED_MASK, max_new_tokens=24, temperature=0.0)
+    preds = []
+    for g in gens:
+        toks = ids_to_tokens(g.detach().cpu().numpy())
+        # cut at EOS
+        toks = toks[: toks.index("</s>")] if "</s>" in toks else toks
+        label_str = bpe_decode_tokens(toks).strip().lower()
+
+        # split multi-label guesses
+        parts = []
+        for sep in [",", ";", "|"]:
+            if sep in label_str:
+                parts = [p.strip() for p in label_str.split(sep) if p.strip()]
+                break
+        if not parts:
+            parts = [label_str] if label_str else []
+
+        # map to canonical hallmarks (no default to 'empty')
+        mapped = []
+        for p in parts:
+            can = to_canonical(p)
+            if can and can not in mapped:
+                mapped.append(can)
+        preds.append(mapped)  # may be []
+    return preds
+
+# --- Run decoding on TEST ---
+model.eval()
+ctx_test = make_context_only(test_df)
+
+B = 32
+preds_all = []
+for i in tqdm(range(0, len(ctx_test), B), desc="Decoding (test)"):
+    batch_ctx = pad_batch(ctx_test[i:i+B]).to(device)
+    preds_all.extend(predict_labels_for_batch(batch_ctx))
+
+# --- Build ground truth (hallmarks only) ---
+y_true = [ normalize_labels(split_labels(s)) for s in test_df["label"].astype(str).tolist() ]
+
+# --- Binarize and score (10 hallmarks only) ---
+from sklearn.metrics import precision_recall_fscore_support
+LABELS = HALLMARKS
+LIDX = {l:i for i,l in enumerate(LABELS)}
+
+def binarize(labs):
+    v = [0]*len(LABELS)
+    for l in labs:
+        if l in LIDX:
+            v[LIDX[l]] = 1
+    return v
+
+Y_true = np.array([binarize(labs) for labs in y_true], dtype=np.int64)
+Y_pred = np.array([binarize(labs) for labs in preds_all], dtype=np.int64)
+
+micro_p, micro_r, micro_f1, _ = precision_recall_fscore_support(Y_true, Y_pred, average='micro', zero_division=0)
+macro_p, macro_r, macro_f1, _ = precision_recall_fscore_support(Y_true, Y_pred, average='macro', zero_division=0)
+
+print(f"\nHALLMARKS-ONLY  Micro  P/R/F1: {micro_p:.4f} / {micro_r:.4f} / {micro_f1:.4f}")
+print(  f"HALLMARKS-ONLY  Macro  P/R/F1: {macro_p:.4f} / {macro_r:.4f} / {macro_f1:.4f}")
+
+perclass = precision_recall_fscore_support(Y_true, Y_pred, average=None, zero_division=0)
+per_df = pd.DataFrame({
+    "label": LABELS,
+    "precision": perclass[0],
+    "recall":    perclass[1],
+    "f1":        perclass[2],
+    "support":   perclass[3],
+}).sort_values("label")
+
+print("\nPer-class results (10 hallmarks):")
+print(per_df.to_string(index=False))
+
+per_df.to_csv("hoc_test_results_per_class.csv", index=False)
+print("Saved: hoc_test_results_per_class.csv")