import json import argparse import os import random import torch import torch.nn as nn from torch.utils.data import Dataset, DataLoader, random_split from transformers import LayoutLMv3TokenizerFast, LayoutLMv3Model from TorchCRF import CRF from torch.optim import AdamW from tqdm import tqdm from sklearn.metrics import precision_recall_fscore_support # --- Configuration --- MAX_BBOX_DIMENSION = 1000 MAX_SHIFT = 30 AUGMENTATION_FACTOR = 1 BASE_MODEL_ID = "microsoft/layoutlmv3-base" # ------------------------- # Step 1: Preprocessing # ------------------------- def preprocess_labelstudio(input_path, output_path): with open(input_path, "r", encoding="utf-8") as f: data = json.load(f) processed = [] print(f"🔄 Starting preprocessing of {len(data)} documents...") for item in data: words = item["data"]["original_words"] bboxes = item["data"]["original_bboxes"] labels = ["O"] * len(words) clamped_bboxes = [] for bbox in bboxes: x_min, y_min, x_max, y_max = bbox new_x_min = max(0, min(x_min, 1000)) new_y_min = max(0, min(y_min, 1000)) new_x_max = max(0, min(x_max, 1000)) new_y_max = max(0, min(y_max, 1000)) if new_x_min > new_x_max: new_x_min = new_x_max if new_y_min > new_y_max: new_y_min = new_y_max clamped_bboxes.append([new_x_min, new_y_min, new_x_max, new_y_max]) if "annotations" in item: for ann in item["annotations"]: for res in ann["result"]: if "value" in res and "labels" in res["value"]: text = res["value"]["text"] tag = res["value"]["labels"][0] text_tokens = text.split() for i in range(len(words) - len(text_tokens) + 1): if words[i:i + len(text_tokens)] == text_tokens: labels[i] = f"B-{tag}" for j in range(1, len(text_tokens)): labels[i + j] = f"I-{tag}" break processed.append({"tokens": words, "labels": labels, "bboxes": clamped_bboxes}) with open(output_path, "w", encoding="utf-8") as f: json.dump(processed, f, indent=2, ensure_ascii=False) return output_path # ------------------------- # Step 1.5: Augmentation # ------------------------- def translate_bbox(bbox, shift_x, shift_y): x_min, y_min, x_max, y_max = bbox new_x_min = max(0, min(x_min + shift_x, 1000)) new_y_min = max(0, min(y_min + shift_y, 1000)) new_x_max = max(0, min(x_max + shift_x, 1000)) new_y_max = max(0, min(y_max + shift_y, 1000)) return [new_x_min, new_y_min, new_x_max, new_y_max] def augment_sample(sample): shift_x = random.randint(-MAX_SHIFT, MAX_SHIFT) shift_y = random.randint(-MAX_SHIFT, MAX_SHIFT) new_sample = sample.copy() new_sample["bboxes"] = [translate_bbox(b, shift_x, shift_y) for b in sample["bboxes"]] return new_sample def augment_and_save_dataset(input_json_path, output_json_path): with open(input_json_path, 'r', encoding="utf-8") as f: training_data = json.load(f) augmented_data = [] for original_sample in training_data: augmented_data.append(original_sample) for _ in range(AUGMENTATION_FACTOR): augmented_data.append(augment_sample(original_sample)) with open(output_json_path, 'w', encoding="utf-8") as f: json.dump(augmented_data, f, indent=2, ensure_ascii=False) return output_json_path # ------------------------- # Step 2: Dataset Class # ------------------------- class LayoutDataset(Dataset): def __init__(self, json_path, tokenizer, label2id, max_len=512): with open(json_path, "r", encoding="utf-8") as f: self.data = json.load(f) self.tokenizer = tokenizer self.label2id = label2id self.max_len = max_len def __len__(self): return len(self.data) def __getitem__(self, idx): item = self.data[idx] words, bboxes, labels = item["tokens"], item["bboxes"], item["labels"] encodings = self.tokenizer(words, boxes=bboxes, padding="max_length", truncation=True, max_length=self.max_len, return_tensors="pt") word_ids = encodings.word_ids(batch_index=0) label_ids = [] for word_id in word_ids: if word_id is None: label_ids.append(self.label2id["O"]) else: label_ids.append(self.label2id.get(labels[word_id], self.label2id["O"])) encodings["labels"] = torch.tensor(label_ids) return {key: val.squeeze(0) for key, val in encodings.items()} # ------------------------- # Step 3: Model Architecture (Non-Linear Head) # ------------------------- class LayoutLMv3CRF(nn.Module): def __init__(self, num_labels): super().__init__() # Initializing from scratch (Base weights only) print(f"🔄 Initializing backbone from {BASE_MODEL_ID}...") self.layoutlm = LayoutLMv3Model.from_pretrained(BASE_MODEL_ID) hidden_size = self.layoutlm.config.hidden_size # NON-LINEAR MLP HEAD # Replacing the simple Linear layer with a deeper architecture self.classifier = nn.Sequential( nn.Linear(hidden_size, hidden_size), nn.GELU(), # Non-linear activation nn.LayerNorm(hidden_size), # Stability for training from scratch nn.Dropout(0.1), nn.Linear(hidden_size, num_labels) ) self.crf = CRF(num_labels) def forward(self, input_ids, bbox, attention_mask, labels=None): outputs = self.layoutlm(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask) sequence_output = outputs.last_hidden_state # Pass through the new non-linear head emissions = self.classifier(sequence_output) if labels is not None: log_likelihood = self.crf(emissions, labels, mask=attention_mask.bool()) return -log_likelihood.mean() else: return self.crf.viterbi_decode(emissions, mask=attention_mask.bool()) # ------------------------- # Step 4: Training + Evaluation # ------------------------- def train_one_epoch(model, dataloader, optimizer, device): model.train() total_loss = 0 for batch in tqdm(dataloader, desc="Training"): batch = {k: v.to(device) for k, v in batch.items()} labels = batch.pop("labels") optimizer.zero_grad() loss = model(**batch, labels=labels) loss.backward() optimizer.step() total_loss += loss.item() return total_loss / len(dataloader) def evaluate(model, dataloader, device, id2label): model.eval() all_preds, all_labels = [], [] with torch.no_grad(): for batch in tqdm(dataloader, desc="Evaluating"): batch = {k: v.to(device) for k, v in batch.items()} labels = batch.pop("labels").cpu().numpy() preds = model(**batch) for p, l, mask in zip(preds, labels, batch["attention_mask"].cpu().numpy()): valid = mask == 1 l_valid = l[valid].tolist() all_labels.extend(l_valid) all_preds.extend(p[:len(l_valid)]) precision, recall, f1, _ = precision_recall_fscore_support(all_labels, all_preds, average="micro", zero_division=0) return precision, recall, f1 # ------------------------- # Step 5: Main Execution # ------------------------- def main(args): labels = ["O", "B-QUESTION", "I-QUESTION", "B-OPTION", "I-OPTION", "B-ANSWER", "I-ANSWER", "B-SECTION_HEADING", "I-SECTION_HEADING", "B-PASSAGE", "I-PASSAGE"] label2id = {l: i for i, l in enumerate(labels)} id2label = {i: l for l, i in label2id.items()} TEMP_DIR = "temp_intermediate_files" os.makedirs(TEMP_DIR, exist_ok=True) # 1. Preprocess & Augment initial_json = os.path.join(TEMP_DIR, "data_bio.json") preprocess_labelstudio(args.input, initial_json) augmented_json = os.path.join(TEMP_DIR, "data_aug.json") final_data_path = augment_and_save_dataset(initial_json, augmented_json) # 2. Setup Data tokenizer = LayoutLMv3TokenizerFast.from_pretrained(BASE_MODEL_ID) dataset = LayoutDataset(final_data_path, tokenizer, label2id, max_len=args.max_len) val_size = int(0.2 * len(dataset)) train_dataset, val_dataset = random_split(dataset, [len(dataset) - val_size, val_size]) train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True) val_loader = DataLoader(val_dataset, batch_size=args.batch_size) # 3. Model device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model = LayoutLMv3CRF(num_labels=len(labels)).to(device) optimizer = AdamW(model.parameters(), lr=args.lr) # 4. Loop for epoch in range(args.epochs): loss = train_one_epoch(model, train_loader, optimizer, device) p, r, f1 = evaluate(model, val_loader, device, id2label) print(f"Epoch {epoch+1} | Loss: {loss:.4f} | F1: {f1:.3f}") ckpt_path = "checkpoints/layoutlmv3_nonlinear_scratch.pth" os.makedirs("checkpoints", exist_ok=True) torch.save(model.state_dict(), ckpt_path) if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--mode", type=str, default="train") parser.add_argument("--input", type=str, required=True) parser.add_argument("--batch_size", type=int, default=4) parser.add_argument("--epochs", type=int, default=10) # Increased for scratch training parser.add_argument("--lr", type=float, default=2e-5) parser.add_argument("--max_len", type=int, default=512) args = parser.parse_args() main(args)