Upload charcnn_bylstm.py

charcnn_bylstm.py

@@ -0,0 +1,730 @@
+# mcq_extractor_updated.py
+import os
+import re
+import io
+import json
+import math
+import pickle
+from collections import Counter, defaultdict
+from typing import List, Tuple
+
+import fitz  # PyMuPDF
+import pytesseract
+from PIL import Image
+import numpy as np
+from tqdm import tqdm
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Dataset, DataLoader
+# use the same CRF you had
+from TorchCRF import CRF  # pip install torchcrf
+
+# ========== CONFIG ==========
+DATA_DIR = "output_data"
+IMAGES_DIR = os.path.join(DATA_DIR, "images")
+os.makedirs(IMAGES_DIR, exist_ok=True)
+PAGE_OCR_CHAR_THRESHOLD = 300
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+MAX_CHAR_LEN = 16
+EMBED_DIM = 100
+CHAR_EMBED_DIM = 30
+CHAR_CNN_OUT = 30
+HIDDEN_SIZE = 256
+BATCH_SIZE = 8
+EPOCHS = 50
+LR = 1e-3
+
+pytesseract.pytesseract.tesseract_cmd = r"D:\prince\New folder\tesseract.exe"
+
+# ========== LABELS (single source of truth) ==========
+LABELS = ["O", "B-QUESTION", "I-QUESTION", "B-OPTION", "I-OPTION", "B-ANSWER", "I-ANSWER",'B-IMAGE','I-IMAGE']
+LABEL2IDX = {l: i for i, l in enumerate(LABELS)}
+IDX2LABEL = {i: l for l, i in LABEL2IDX.items()}
+
+
+# ---------- small utility classes ----------
+class Vocab:
+    def __init__(self, min_freq=1, unk_token="<UNK>", pad_token="<PAD>"):
+        self.min_freq = min_freq
+        self.unk_token = unk_token
+        self.pad_token = pad_token
+        self.freq = Counter()
+        self.itos = []
+        self.stoi = {}
+
+    def add_sentence(self, toks):
+        self.freq.update(toks)
+
+    def build(self):
+        items = [tok for tok, c in self.freq.items() if c >= self.min_freq]
+        items = [self.pad_token, self.unk_token] + sorted(items)
+        self.itos = items
+        self.stoi = {s: i for i, s in enumerate(self.itos)}
+
+    def __len__(self):
+        return len(self.itos)
+
+
+# ========== PDF / tokenization utils (keep yours, slightly cleaned) ==========
+def clean_text_token(t):
+    """Normalizes special characters in a token."""
+    return t.replace("\u2011", "-")  # normalize hyphen
+
+
+PAGE_OCR_CHAR_THRESHOLD = 50
+
+def extract_pdf_pages(path: str):
+    """
+    Extracts content from PDF pages.
+    Returns a list of pages with:
+      - 'width', 'height' -> page dimensions
+      - 'blocks' -> text blocks with bbox
+      - 'images' -> images with bbox and PIL image
+    """
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"The file was not found: {path}")
+
+    doc = fitz.open(path)
+    pages = []
+
+    for pno, page in enumerate(doc):
+        w, h = page.rect.width, page.rect.height
+
+        # Extract text blocks
+        raw_blocks = page.get_text("blocks", sort=True)
+        text_blocks = []
+        for b in raw_blocks:
+            x0, y0, x1, y1, text, block_no, block_type = b
+            if block_type != 0:  # 0 = text block
+                continue
+            text = text.strip().replace("\n", " ")
+            if text:
+                text_blocks.append({
+                    "bbox": (x0, y0, x1, y1),
+                    "text": text,
+                    "font_size": None  # can optionally extract from span if needed
+                })
+
+        # Extract images
+        images = []
+        for img_info in page.get_images(full=True):
+            xref = img_info[0]
+            try:
+                base_image = doc.extract_image(xref)
+                img_bytes = base_image["image"]
+                img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
+                img_rect = page.get_image_bbox(img_info)
+                images.append({"bbox": (img_rect.x0, img_rect.y0, img_rect.x1, img_rect.y1), "image": img})
+            except Exception as e:
+                print(f"Warning: Could not extract image {xref} on page {pno+1}. Error: {e}")
+
+        # OCR fallback if text is too little
+        total_chars = sum(len(b["text"]) for b in text_blocks)
+        if total_chars < PAGE_OCR_CHAR_THRESHOLD:
+            pix = page.get_pixmap(dpi=300)
+            img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
+            ocr_text = pytesseract.image_to_string(img)
+            if ocr_text.strip():
+                text_blocks = [{"bbox": (0, 0, w, h), "text": ocr_text.strip(), "font_size": None}]
+
+        pages.append({"width": w, "height": h, "blocks": text_blocks, "images": images})
+
+    doc.close()
+    return pages
+
+
+import re
+
+IMAGES_DIR = "output_data/images"
+
+def split_blocks_into_tokens(pages):
+    """
+    Tokenizes text blocks and sorts them based on page layout (single or two-column).
+    Returns a list of pages, each containing a list of token dicts.
+    """
+    token_re = re.compile(r"\w+|[^\w\s]", re.UNICODE)
+    all_pages_tokens = []
+
+    for pidx, page in enumerate(pages):
+        tokens = []
+        page_w, page_h = page["width"], page["height"]
+        mid_x = page_w / 2
+
+        # Detect if page is two-column
+        left_count, right_count, spanning_count = 0, 0, 0
+        gutter = 0.1 * page_w
+        for b in page["blocks"]:
+            x0, y0, x1, y1 = b["bbox"]
+            if y0 < 0.05 * page_h or y1 > 0.95 * page_h:  # ignore headers/footers
+                continue
+            if x1 < mid_x - gutter/2:
+                left_count += 1
+            elif x0 > mid_x + gutter/2:
+                right_count += 1
+            elif x0 < mid_x and x1 > mid_x:
+                spanning_count += 1
+        is_two_column = left_count > 3 and right_count > 3 and spanning_count <= 2
+
+        # Tokenize blocks
+        for bidx, block in enumerate(page["blocks"]):
+            x0, y0, x1, y1 = block["bbox"]
+            text = block["text"].replace("\u00ad", "")
+            toks = token_re.findall(text)
+            if not toks:
+                continue
+            total_chars = sum(len(t) for t in toks)
+            cur_x = x0
+            for tok in toks:
+                tok_width = (len(tok)/total_chars)* (x1 - x0) if total_chars>0 else (x1-x0)/len(toks)
+                tokens.append({
+                    "text": clean_text_token(tok),
+                    "x0": cur_x, "y0": y0,
+                    "x1": cur_x + tok_width, "y1": y1,
+                    "font_size": block.get("font_size"),
+                    "page_no": pidx+1,
+                    "block_idx": bidx
+                })
+                cur_x += tok_width
+
+        # Sort tokens based on layout
+        if is_two_column:
+            tokens.sort(key=lambda t: (0 if t['x0'] < mid_x else 1, t['y0'], t['x0']))
+        else:
+            tokens.sort(key=lambda t: (t['y0'], t['x0']))
+
+        all_pages_tokens.append(tokens)
+    return all_pages_tokens
+
+
+def assign_images_to_tokens(pages, all_pages_tokens):
+    """
+    Inserts image placeholders into the token stream.
+    """
+    if not os.path.exists(IMAGES_DIR):
+        os.makedirs(IMAGES_DIR)
+
+    for pidx, page in enumerate(pages):
+        tokens = all_pages_tokens[pidx]
+        for img_idx, imrec in enumerate(page["images"]):
+            img_name = f"page{pidx+1}_img{img_idx+1}.png"
+            imrec["image"].save(os.path.join(IMAGES_DIR, img_name))
+            img_center_y = (imrec["bbox"][1]+imrec["bbox"][3])/2
+            if not tokens:
+                insert_idx = 0
+            else:
+                closest_token = min(tokens, key=lambda t: abs((t["y0"]+t["y1"])/2 - img_center_y))
+                insert_idx = tokens.index(closest_token)+1
+            tokens.insert(insert_idx, {
+                "text": f"[IMAGE: {img_name}]",
+                "x0": imrec["bbox"][0], "y0": imrec["bbox"][1],
+                "x1": imrec["bbox"][2], "y1": imrec["bbox"][3],
+                "font_size": None,
+                "page_no": pidx+1,
+                "block_idx": -1,
+                "is_image": True
+            })
+        all_pages_tokens[pidx] = tokens
+    return all_pages_tokens
+
+
+
+
+# ========== Dataset ==========
+def orthographic_features(token_text):
+    return [
+        int(token_text[0].isupper()) if token_text and token_text[0].isalpha() else 0,
+        int(token_text.isupper()),
+        int(any(ch.isdigit() for ch in token_text)),
+        int(len(token_text) == 1 and re.match(r'\W', token_text) is not None)
+    ]
+
+class MCQTokenDataset(Dataset):
+    def __init__(self, pages_tokens, word_vocab, char_vocab, labels_per_token=None):
+        self.samples = []
+        self.labels = []
+
+        if labels_per_token:
+            for toks, lbls in zip(pages_tokens, labels_per_token):
+                if len(toks) == 0:
+                    continue  # skip empty pages
+                if len(toks) != len(lbls):
+                    raise ValueError(f"Token/label length mismatch: {len(toks)} vs {len(lbls)}")
+                self.samples.append(toks)
+                self.labels.append(lbls)
+        else:
+            self.samples = [p for p in pages_tokens if len(p) > 0]
+        self.word_vocab = word_vocab
+        self.char_vocab = char_vocab
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, idx):
+        toks = self.samples[idx]
+
+        # ✅ Make sure every token has text
+        words = []
+        safe_toks = []
+        for t in toks:
+            if isinstance(t, dict) and "text" in t:
+                txt = t["text"]
+                safe_toks.append(t)
+            elif isinstance(t, str):
+                txt = t
+                safe_toks.append({"text": txt, "x0": 0, "x1": 0, "y0": 0, "y1": 0, "font_size": 0.0})
+            else:
+                txt = str(t)
+                safe_toks.append({"text": txt, "x0": 0, "x1": 0, "y0": 0, "y1": 0, "font_size": 0.0})
+            words.append(txt)
+
+        toks = safe_toks  # Use normalized tokens downstream
+
+        word_ids = [self.word_vocab.stoi.get(w, self.word_vocab.stoi[self.word_vocab.unk_token]) for w in words]
+
+        char_ids = []
+        for w in words:
+            chs = [self.char_vocab.stoi.get(ch, self.char_vocab.stoi[self.char_vocab.unk_token]) for ch in
+                   w[:MAX_CHAR_LEN]]
+            if len(chs) < MAX_CHAR_LEN:
+                chs += [self.char_vocab.stoi[self.char_vocab.pad_token]] * (MAX_CHAR_LEN - len(chs))
+            char_ids.append(chs)
+
+        x_centers = [(t["x0"] + t["x1"]) / 2.0 for t in toks]
+        y_centers = [(t["y0"] + t["y1"]) / 2.0 for t in toks]
+        max_x = max([t["x1"] for t in toks]) if toks else 1.0
+        max_y = max([t["y1"] for t in toks]) if toks else 1.0
+
+        if max_x == 0:
+            max_x = 1.0
+        if max_y == 0:
+            max_y = 1.0
+
+        x_norm = [xc / max_x for xc in x_centers]
+        y_norm = [yc / max_y for yc in y_centers]
+
+        font_sizes = [float(t.get("font_size") or 0.0) for t in toks]
+        ortho_feats = [orthographic_features(w) for w in words]
+
+        labels = None
+        if self.labels:
+            lbls = self.labels[idx]
+            labels = [LABEL2IDX[l] for l in lbls]
+
+        return {
+            "word_ids": torch.LongTensor(word_ids),
+            "char_ids": torch.LongTensor(char_ids),
+            "x_norm": torch.FloatTensor(x_norm),
+            "y_norm": torch.FloatTensor(y_norm),
+            "font_sizes": torch.FloatTensor(font_sizes),
+            "ortho": torch.FloatTensor(ortho_feats),
+            "labels": torch.LongTensor(labels) if labels is not None else None,
+            "tokens": toks
+        }
+
+
+def collate_batch(batch):
+    batch = [item for item in batch if item["word_ids"].size(0) > 0]  # remove empty sequences
+    if len(batch) == 0:
+        return None  # or raise error
+
+    max_len = max(item["word_ids"].size(0) for item in batch)
+    batch_size = len(batch)
+
+    word_pad = torch.zeros((batch_size, max_len), dtype=torch.long)
+    char_pad = torch.zeros((batch_size, max_len, MAX_CHAR_LEN), dtype=torch.long)
+    x_pad = torch.zeros((batch_size, max_len), dtype=torch.float)
+    y_pad = torch.zeros((batch_size, max_len), dtype=torch.float)
+    font_pad = torch.zeros((batch_size, max_len), dtype=torch.float)
+    ortho_pad = torch.zeros((batch_size, max_len, 4), dtype=torch.float)
+    mask = torch.zeros((batch_size, max_len), dtype=torch.bool)
+    label_pad = torch.full((batch_size, max_len), LABEL2IDX["O"], dtype=torch.long)  # use O as default
+    tokens_list = []
+
+    for i, item in enumerate(batch):
+        L = item["word_ids"].size(0)
+        word_pad[i, :L] = item["word_ids"]
+        char_pad[i, :L, :] = item["char_ids"]
+        x_pad[i, :L] = item["x_norm"]
+        y_pad[i, :L] = item["y_norm"]
+        font_pad[i, :L] = item["font_sizes"]
+        ortho_pad[i, :L, :] = item["ortho"]
+        mask[i, :L] = 1
+        if item["labels"] is not None and item["labels"].size(0) == L:
+            label_pad[i, :L] = item["labels"]
+        tokens_list.append(item["tokens"])
+
+    return {
+        "words": word_pad,
+        "chars": char_pad,
+        "x": x_pad,
+        "y": y_pad,
+        "font": font_pad,
+        "ortho": ortho_pad,
+        "mask": mask,
+        "labels": label_pad,
+        "tokens": tokens_list
+    }
+
+# ========== MODEL ==========
+class CharCNNEncoder(nn.Module):
+    def __init__(self, char_vocab_size, char_emb_dim, out_dim, kernel_sizes=(3,4,5)):
+        super().__init__()
+        self.char_emb = nn.Embedding(char_vocab_size, char_emb_dim, padding_idx=0)
+        convs = []
+        for k in kernel_sizes:
+            convs.append(nn.Conv1d(char_emb_dim, out_dim, kernel_size=k))
+        self.convs = nn.ModuleList(convs)
+        self.out_dim = out_dim * len(convs)
+
+    def forward(self, char_ids):
+        B, L, C = char_ids.size()
+        emb = self.char_emb(char_ids.view(B * L, C))
+        emb = emb.transpose(1,2)
+        outs = []
+        for conv in self.convs:
+            c = conv(emb)
+            c = torch.relu(c)
+            c = torch.max(c, dim=2)[0]
+            outs.append(c)
+        res = torch.cat(outs, dim=1)
+        return res.view(B, L, -1)
+
+class MCQTagger(nn.Module):
+    def __init__(self, vocab_size, char_vocab_size, n_labels):
+        super().__init__()
+        self.word_emb = nn.Embedding(vocab_size, EMBED_DIM, padding_idx=0)
+        self.char_enc = CharCNNEncoder(char_vocab_size, CHAR_EMBED_DIM, CHAR_CNN_OUT)
+        in_dim = EMBED_DIM + self.char_enc.out_dim + 2 + 1 + 4
+        self.bilstm = nn.LSTM(in_dim, HIDDEN_SIZE // 2, num_layers=1, batch_first=True, bidirectional=True)
+        self.ff = nn.Linear(HIDDEN_SIZE, n_labels)
+        self.crf = CRF(n_labels, batch_first=True)
+
+    def forward_emissions(self, words, chars, x, y, font, ortho, mask):
+        # return raw emissions (before CRF) so we can obtain per-token probs
+        wemb = self.word_emb(words)
+        cenc = self.char_enc(chars)
+        numeric = torch.cat([x.unsqueeze(-1), y.unsqueeze(-1), font.unsqueeze(-1), ortho], dim=-1)
+        enc_in = torch.cat([wemb, cenc, numeric], dim=-1)
+        packed_out, _ = self.bilstm(enc_in)
+        emissions = self.ff(packed_out)
+        return emissions
+
+    def forward(self, words, chars, x, y, font, ortho, mask, labels=None, class_weights=None, alpha=0.7):
+        emissions = self.forward_emissions(words, chars, x, y, font, ortho, mask)
+        if labels is not None:
+            crf_loss = -self.crf(emissions, labels, mask=mask, reduction='mean')
+            if class_weights is not None:
+                ce_loss_fn = nn.CrossEntropyLoss(weight=class_weights.to(emissions.device), ignore_index=-1)
+                ce_loss = ce_loss_fn(emissions.view(-1, emissions.size(-1)), labels.view(-1))
+                loss = alpha * crf_loss + (1 - alpha) * ce_loss
+            else:
+                loss = crf_loss
+            return loss
+        else:
+            pred = self.crf.decode(emissions, mask=mask)
+            return pred
+
+# helper: get softmax probs per token from emissions
+def emissions_to_probs(emissions, mask):
+    # emissions: (B, L, C)
+    probs = F.softmax(emissions, dim=-1)  # (B,L,C)
+    probs = probs.cpu().numpy()
+    masks = mask.cpu().numpy()
+    # return as list of arrays per example (only active tokens)
+    out = []
+    for i in range(probs.shape[0]):
+        L = masks[i].sum()
+        out.append(probs[i][:L])
+    return out
+
+# ========== training/eval ==========
+def compute_class_weights(labels_list, num_labels):
+    all_labels_flat = [lbl for page in labels_list for lbl in page]
+    counts = Counter(all_labels_flat)
+    total = sum(counts.values())
+    weights = []
+    for i in range(num_labels):
+        count = counts.get(i, 0)
+        if count == 0:
+            w = 1.0
+        else:
+            w = total / (num_labels * count)
+        if IDX2LABEL[i] in ["B-QUESTION", "B-OPTION"]:
+            w *= 2.0
+        weights.append(w)
+    return torch.tensor(weights, dtype=torch.float)
+
+def eval_model(model, data_loader):
+    model.eval()
+    all_true = []
+    all_pred = []
+    with torch.no_grad():
+        for batch in tqdm(data_loader, desc="Eval"):
+            words = batch["words"].to(DEVICE)
+            chars = batch["chars"].to(DEVICE)
+            x = batch["x"].to(DEVICE)
+            y = batch["y"].to(DEVICE)
+            font = batch["font"].to(DEVICE)
+            ortho = batch["ortho"].to(DEVICE)
+            mask = batch["mask"].to(DEVICE)
+            labels = batch["labels"].to(DEVICE)
+            preds = model(words, chars, x, y, font, ortho, mask, labels=None)
+            for i in range(len(preds)):
+                L = mask[i].sum().item()
+                pred_seq = preds[i][:L]
+                true_seq = labels[i][:L].cpu().numpy().tolist()
+                all_pred.extend(pred_seq)
+                all_true.extend(true_seq)
+    # compute token-level micro F1 excluding O maybe; here we compute micro across all labels
+    from sklearn.metrics import precision_recall_fscore_support
+    p, r, f1, _ = precision_recall_fscore_support(all_true, all_pred, average='micro', zero_division=0)
+    return p, r, f1
+
+def train_model(model, train_loader, val_loader, epochs=EPOCHS, class_weights=None):
+    model.to(DEVICE)
+    optim = torch.optim.Adam(model.parameters(), lr=LR)
+    best_val_f1 = 0.0
+    for ep in range(1, epochs+1):
+        model.train()
+        running_loss = 0.0
+        for batch in tqdm(train_loader, desc=f"Train E{ep}"):
+            optim.zero_grad()
+            words = batch["words"].to(DEVICE)
+            chars = batch["chars"].to(DEVICE)
+            x = batch["x"].to(DEVICE)
+            y = batch["y"].to(DEVICE)
+            font = batch["font"].to(DEVICE)
+            ortho = batch["ortho"].to(DEVICE)
+            mask = batch["mask"].to(DEVICE)
+            labels = batch["labels"].to(DEVICE)
+            loss = model(words, chars, x, y, font, ortho, mask, labels, class_weights=class_weights)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optim.step()
+            running_loss += loss.item()
+        avg_loss = running_loss / max(1, len(train_loader))
+        print(f"Epoch {ep} train loss {avg_loss:.4f}")
+        p, r, f1 = eval_model(model, val_loader)
+        print(f"VAL p={p:.4f} r={r:.4f} f1={f1:.4f}")
+        if f1 > best_val_f1:
+            best_val_f1 = f1
+            torch.save(model.state_dict(), os.path.join(DATA_DIR, "best_mcq_tagger.pt"))
+    print("Training complete. Best val F1:", best_val_f1)
+    return model
+
+# ========== helpers to save/load vocabs ==========
+def build_vocabs(pages_tokens):
+    word_vocab = Vocab(min_freq=1)
+    char_vocab = Vocab(min_freq=1, unk_token="<CUNK>", pad_token="<CPAD>")
+
+    for p in pages_tokens:
+        for tok in p:
+            # ✅ Always convert to string safely
+            if isinstance(tok, dict) and "text" in tok:
+                text_value = tok["text"]
+            elif isinstance(tok, str):
+                text_value = tok
+            else:
+                text_value = str(tok)
+
+            word_vocab.add_sentence([text_value])
+            for ch in text_value[:MAX_CHAR_LEN]:
+                char_vocab.add_sentence([ch])
+
+    word_vocab.build()
+    char_vocab.build()
+    return word_vocab, char_vocab
+
+def save_vocabs(path, word_vocab, char_vocab):
+    with open(path, "wb") as f:
+        pickle.dump((word_vocab, char_vocab), f)
+
+def load_vocabs(path):
+    with open(path, "rb") as f:
+        return pickle.load(f)
+
+# ========== reconstruction (unchanged) ==========
+def reconstruct_mcqs_from_tokens(tokens, preds):
+    mcqs = []
+    i = 0
+    N = len(tokens)
+    fragments = []
+    while i < N:
+        label = IDX2LABEL[preds[i]]
+        if label.startswith("B-QUESTION"):
+            if fragments and "question" in fragments[-1]:
+                mcqs.append(fragments[-1])
+            q_toks = [tokens[i]["text"]]
+            i += 1
+            while i < N and IDX2LABEL[preds[i]].startswith("I-QUESTION"):
+                q_toks.append(tokens[i]["text"])
+                i += 1
+            fragments.append({"question": " ".join(q_toks), "options": [], "answer": None})
+        elif fragments:
+            lab = IDX2LABEL[preds[i]]
+            if lab.startswith("B-OPTION"):
+                otoks = [tokens[i]["text"]]
+                i += 1
+                while i < N and IDX2LABEL[preds[i]].startswith("I-OPTION"):
+                    otoks.append(tokens[i]["text"])
+                    i += 1
+                fragments[-1]["options"].append(" ".join(otoks))
+            elif lab.startswith("B-ANSWER"):
+                atoks = [tokens[i]["text"]]
+                i += 1
+                while i < N and IDX2LABEL[preds[i]].startswith("I-ANSWER"):
+                    atoks.append(tokens[i]["text"])
+                    i += 1
+                fragments[-1]["answer"] = " ".join(atoks)
+            else:
+                i += 1
+        else:
+            i += 1
+
+    if fragments and "question" in fragments[-1]:
+        mcqs.append(fragments[-1])
+
+    # ✅ filter only "perfect" mcqs: must have a question and at least one option
+    mcqs = [m for m in mcqs if m.get("question") and m.get("options")]
+
+    return mcqs
+
+def convert_labels_to_indices(all_labels):
+    all_labels_indices = [
+        [LABEL2IDX[l] for l in page] for page in all_labels
+    ]
+    return all_labels_indices
+def demo_inference(pdf_path, model_path, vocab_path):
+    import json
+    from torch.utils.data import DataLoader
+
+    # Load vocabs
+    word_vocab, char_vocab = load_vocabs(vocab_path)
+
+    # Load model
+    model = MCQTagger(len(word_vocab), len(char_vocab), n_labels=len(LABELS))
+    model.load_state_dict(torch.load(model_path, map_location=DEVICE))
+    model.to(DEVICE)
+    model.eval()
+
+    # Extract + tokenize PDF
+    pages = extract_pdf_pages(pdf_path)
+    pages_tokens = split_blocks_into_tokens(pages)
+    pages_tokens = assign_images_to_tokens(pages, pages_tokens)
+
+    # Dataset + loader
+    dataset = MCQTokenDataset(pages_tokens, word_vocab, char_vocab, labels_per_token=None)
+    loader = DataLoader(dataset, batch_size=1, collate_fn=collate_batch)
+
+    all_mcqs = []
+    all_preds = []
+    with torch.no_grad():
+        for batch in loader:
+            words = batch["words"].to(DEVICE)
+            chars = batch["chars"].to(DEVICE)
+            x = batch["x"].to(DEVICE)
+            y = batch["y"].to(DEVICE)
+            font = batch["font"].to(DEVICE)
+            ortho = batch["ortho"].to(DEVICE)
+            mask = batch["mask"].to(DEVICE)
+            tokens = batch["tokens"][0]
+
+            preds = model(words, chars, x, y, font, ortho, mask, labels=None)
+            preds = preds[0]  # batch size = 1
+            all_preds.append(preds)
+
+            mcqs = reconstruct_mcqs_from_tokens(tokens, preds)
+            all_mcqs.extend(mcqs)
+
+    # Save to JSON (optional)
+    out_path = os.path.join(DATA_DIR, f"cnn_{os.path.basename(pdf_path)}.json")
+    with open(out_path, "w", encoding="utf-8") as f:
+        json.dump(all_mcqs, f, ensure_ascii=False, indent=2)
+
+    print(f"✅ Results saved to {out_path}")
+    return all_mcqs, all_preds
+
+# if run as script, keep legacy demo functions etc. (omitted for brevity)
+if __name__ == "__main__":
+
+
+        #("augmented_data/english_21-50_labels_augmented_1.json", "augmented_data/english_21-50_labels_augmented_1.json"),
+        #("augmented_data/english_21-50_tokens_augmented_2.json", "augmented_data/english_21-50_labels_augmented_2.json"),
+        #("augmented_data/english_21-50_tokens_augmented_3.json", "augmented_data/english_21-50_labels_augmented_3.json"),
+        #("augmented_data/english_21-50_tokens_augmented_4.json", "augmented_data/english_21-50_labels_augmented_4.json"),
+        #("augmented_data/english_21-50_tokens_augmented_5.json", "augmented_data/english_21-50_labels_augmented_5.json")
+    with open("merged_tokens_labels.json", "r", encoding="utf-8") as f:
+        merged_data = json.load(f)
+    all_pages_tokens=[]
+    all_labels=[]
+        # group by page if needed — assuming all tokens are from one page,
+        # otherwise you can group by "page_no"
+    from itertools import groupby
+
+    merged_data.sort(key=lambda x: x.get("page_no", 0))
+    pages = []
+    for page_no, group in groupby(merged_data, key=lambda x: x.get("page_no", 0)):
+            group = list(group)
+            tokens = []
+            labels = []
+            for item in group:
+                tokens.append({
+                    "text": item.get("text", ""),
+                    "x0": item.get("x0", 0),
+                    "y0": item.get("y0", 0),
+                    "x1": item.get("x1", 0),
+                    "y1": item.get("y1", 0),
+                    "font_size": item.get("font_size", 0),
+                    "page_no": item.get("page_no", 0),
+                    "block_idx": item.get("block_idx", 0)
+                })
+                labels.append(item.get("label", "O"))
+            all_pages_tokens.append(tokens)
+            all_labels.append(labels)
+
+    # 🔀 Split into training and validation
+    split_idx = int(len(all_pages_tokens) * 0.8)
+    train_pages_tokens = all_pages_tokens[:split_idx]
+    train_labels = all_labels[:split_idx]
+    val_pages_tokens = all_pages_tokens[split_idx:]
+    val_labels = all_labels[split_idx:]
+
+    print(f"Training on {len(train_labels)} pages, validating on {len(val_labels)} pages")
+
+    # 🧮 Compute class weights
+    all_labels_indices = convert_labels_to_indices(all_labels)
+    class_weights = compute_class_weights(all_labels_indices, len(LABELS)).to(DEVICE)
+    print("Class weights:", class_weights)
+
+
+    # 🏗️ Build vocabularies
+    word_vocab, char_vocab = build_vocabs(train_pages_tokens)
+
+    # 📦 Build datasets
+    dataset_train = MCQTokenDataset(train_pages_tokens, word_vocab, char_vocab, labels_per_token=train_labels)
+    dataset_val = MCQTokenDataset(val_pages_tokens, word_vocab, char_vocab, labels_per_token=val_labels)
+
+    # 🔄 Data loaders
+    train_loader = DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True, collate_fn=collate_batch)
+    val_loader = DataLoader(dataset_val, batch_size=BATCH_SIZE, collate_fn=collate_batch)
+
+    # 🧠 Train model
+    model = MCQTagger(len(word_vocab), len(char_vocab), len(LABELS))
+    train_model(model, train_loader, val_loader, epochs=EPOCHS, class_weights=class_weights)
+
+    # 💾 Save vocabs for later inference
+    os.makedirs(DATA_DIR, exist_ok=True)
+    with open(os.path.join(DATA_DIR, "vocabs.pkl"), "wb") as f:
+        pickle.dump((word_vocab, char_vocab), f)
+    train_loader = DataLoader(dataset_train, batch_size=2, shuffle=True, collate_fn=collate_batch)
+
+    # Debug: check if rare labels appear in a batch
+    for batch in train_loader:
+        labels_in_batch = batch['labels']  # adjust key based on your dataset collate
+        unique_labels = torch.unique(torch.cat([torch.tensor([0, 1]), torch.tensor([2, 3])]))
+        print("Labels in batch:", unique_labels)
+        break
+    print("✅ Training finished. Model + vocabs saved.")
+