Upload 20_entities_top_50_25's state dict

.gitattributes

@@ -83,3 +83,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 12_13_lr_16/logs/12_13_lr_16_log_plot.jpg filter=lfs diff=lfs merge=lfs -text
 8_entities_top_span_self_ensemble_no_weight_20/logs/8_entities_top_span_self_ensemble_no_weight_20_log_plot.jpg filter=lfs diff=lfs merge=lfs -text
 20_weight_5_21/logs/20_weight_5_21_log_plot.jpg filter=lfs diff=lfs merge=lfs -text
+20_entities_top_50_25/logs/20_entities_top_50_25_log_plot.jpg filter=lfs diff=lfs merge=lfs -text

20_entities_top_50_25/20_entities_top_50_25.py

@@ -0,0 +1,2419 @@
+# %% [code]
+get_ipython().system('pip install evaluate seqeval underthesea positional-encodings[pytorch] pytorch-crf')
+
+# %% [code]
+import warnings
+warnings.filterwarnings('ignore')
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, TensorDataset, DataLoader
+import torch.nn.functional as F
+import albumentations as albu
+from transformers import AutoTokenizer, AutoModel
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+from positional_encodings.torch_encodings import PositionalEncoding1D
+from torchcrf import CRF
+
+from sklearn.metrics import f1_score
+from sklearn.preprocessing import MinMaxScaler, StandardScaler
+from scipy.spatial.transform import Rotation as R
+from sklearn.model_selection import KFold, StratifiedGroupKFold, GroupKFold, StratifiedKFold
+from sklearn.metrics import precision_recall_fscore_support
+from timm.utils import ModelEmaV3
+import timm
+
+import os
+import gc
+import json
+from pathlib import Path
+import pickle
+from tqdm.auto import tqdm
+import copy
+import numpy as np
+import pandas as pd
+import polars as pl
+from PIL import Image
+import time
+from tqdm import tqdm
+from matplotlib import pyplot as plt
+import seaborn as sns
+from multiprocessing import Manager as MemoryManager
+from functools import lru_cache
+import shutil
+import glob
+import cv2
+import random
+import re
+import joblib
+import math
+from huggingface_hub import HfApi, snapshot_download
+import evaluate
+from underthesea import word_tokenize as vi_tokenize_tool
+import spacy
+en_tokenize_tool = spacy.load("en_core_web_sm")
+from collections import defaultdict, Counter
+
+# %% [code]
+# Global config
+SEEDS = [26092004]
+topk = 1
+nfolds = 5
+only_fold_idx = 0
+test_only = 0
+debug_only = 0
+
+# Config thư mục
+dataset = 'kltn/only_entities' # conll003, ontonotes, phoner, vietbio, vietmed, vimed, kltn/only_entities, kltn/raw
+root_dir = f'/kaggle/input/notebooks/sambui22022517/kltn-data/{dataset}' ## Thư mục chứa file train, val, test
+train_dir = f'{root_dir}'
+# val_dir = f'{root_dir}/val'
+test_dir = f'{root_dir}'
+
+# Config checkpoints
+
+# Config training
+epochs = 18 if not debug_only else 2
+batch_size = 32
+device = "cuda" if torch.cuda.is_available() else "cpu"
+# # Thêm biến toàn cục nào đó vào đây
+repo_name = 'SS3M/kltn-experiments'
+state_dict_save_name = "20_entities_top_50_25"
+checkpoints_dir = state_dict_save_name
+pretrained_dir = "/kaggle/working"
+os.makedirs(f'{checkpoints_dir}', exist_ok=True)
+
+backbone_model_name = "bert-base-uncased" if dataset in ["conll003", "ontonotes"] else "vinai/phobert-base"
+word_tokenize = lambda text: [token.text for token in en_tokenize_tool(text)] if dataset == dataset in ["conll003", "ontonotes"] else vi_tokenize_tool(text)
+max_len_dict = {
+    'kltn/raw': 256, 
+    'kltn/only_entities': 68, 
+    'conll003': 46, 
+    'ontonotes': 61, 
+    'phoner': 68, 
+    'vietbio': 125, 
+    'vietmed': 36, 
+    'vimed': 100, 
+}
+zero_entities_rate_dict = {
+    'kltn/raw': 1000, 
+    'kltn/only_entities': 0.2, 
+    'conll003': 1000, # mean keep all zero-entities samples
+    'ontonotes': 1000, 
+    'phoner': 1000, 
+    'vietbio': 1000, 
+    'vietmed': 1000, 
+    'vimed': 1000, 
+}
+
+max_len = max_len_dict[dataset]
+max_n_parts = 3 if dataset in ['kltn/raw'] else 1
+max_span_len = 10
+zero_entities_rate = zero_entities_rate_dict[dataset]
+
+# Trainer
+trainer_params = {
+    "training_time": "00:11:30:00",
+    "eval_mode": "max",
+    "topk": topk,
+    "save_name": state_dict_save_name,
+    "save_best": True,
+    "save_last": True,
+    "device": device, 
+    "logging": True, 
+    "logging_file": True, 
+    "checkpoints_dir": checkpoints_dir, 
+    "early_stopping": 30, 
+    "eval_from_ratio": 0.4, 
+    "eval_every": 1, 
+    "schedule_in_step": False, 
+    "use_ema": True, 
+    "ema_from_ratio": 0.3, 
+    "ema_decay": 0.9995, 
+    "max_grad_norm": 200.0, 
+    "return_best": True, 
+    "return_last": True, 
+}
+
+# Memory
+train_memory_params = {
+    'max_len': max_len, 
+    'max_n_parts': max_n_parts, 
+    'max_span_len': max_span_len, 
+    'weight_rate': None, 
+}
+val_memory_params = {
+    'max_len': max_len, 
+    'max_n_parts': max_n_parts, 
+    'max_span_len': max_span_len, 
+    'weight_rate': None, 
+}
+
+# Data Loader
+def seed_worker(worker_id):
+    worker_seed = torch.initial_seed() % 2**32
+    np.random.seed(worker_seed)
+    random.seed(worker_seed)
+    
+train_loader_params = {
+    'batch_size': batch_size, 
+    'shuffle': True,
+    'pin_memory':True, 
+    'num_workers': 2, 
+    'drop_last': False, 
+    'worker_init_fn': seed_worker, 
+    'persistent_workers': False, 
+}
+val_loader_params = {
+    'batch_size': batch_size, 
+    'shuffle': False,
+    'pin_memory':True, 
+    'num_workers': 1, 
+    'drop_last': False, 
+    'worker_init_fn': seed_worker, 
+    'persistent_workers': False, 
+}
+
+# Model
+model_params = {
+    'backbone_model_name': backbone_model_name, 
+    'max_span_len': max_span_len, 
+    'topk_spans': 50, 
+    'keep_neighbor': 2, 
+}
+
+# Loss Func
+loss_func_params = {
+    'lambda_ce': 1.0,
+}
+eval_func_params = {}
+
+# Optim
+optim_params = {
+    'name': 'AdamW', 
+    'lr': 1e-4,
+    'weight_decay': 1e-4,
+}
+scheduler_params = {
+    'name': 'CosineAnnealingLR',
+    'T_max': 20,  # Số epoch để hoàn thành một chu kỳ giảm LR
+    'eta_min': 1e-6  # Learning rate nhỏ nhất trong chu kỳ
+}
+
+# %% [code]
+def set_seed(seed=42):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)  # if using multi-GPU
+    torch.use_deterministic_algorithms(False)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+    os.environ['PYTHONHASHSEED'] = str(seed)
+
+# %% [code]
+class CustomLoss(nn.Module):
+    def __init__(self, lambda_ce=1.0):
+        super().__init__()
+        self.lambda_ce = lambda_ce
+
+        self.ce = nn.CrossEntropyLoss(
+            ignore_index=-100,
+            reduction='none'
+        )
+
+    def forward(
+        self,
+        logits, labels, weights,   # weights: (B, N)
+        start_logits, start_labels,
+        end_logits, end_labels,
+    ):
+        # =====================================
+        # SPAN LOSS
+        # =====================================
+
+        B, N, C = logits.shape
+
+        flat_logits = logits.reshape(-1, C)
+        flat_labels = labels.reshape(-1)
+        flat_weights = weights.reshape(-1)
+
+        valid_mask = flat_labels != -100
+
+        if valid_mask.any():
+            ce_loss = self.ce(flat_logits, flat_labels)  # (B*N,)
+
+            ce_loss = ce_loss[valid_mask]
+            valid_weights = flat_weights[valid_mask]
+
+            loss = (ce_loss * valid_weights).sum() / valid_weights.sum().clamp(min=1e-8)
+
+        else:
+            loss = logits.sum() * 0.0
+
+        # =====================================
+        # START LOSS
+        # =====================================
+
+        B, L, C = start_logits.shape
+
+        start_logits_flat = start_logits.reshape(B * L, C)
+        start_labels_flat = start_labels.reshape(-1)
+
+        start_loss = self.ce(start_logits_flat, start_labels_flat)
+        start_valid = start_labels_flat != -100
+
+        if start_valid.any():
+            start_loss = start_loss[start_valid].mean()
+        else:
+            start_loss = logits.sum() * 0.0
+
+        # =====================================
+        # END LOSS
+        # =====================================
+
+        end_logits_flat = end_logits.reshape(B * L, C)
+        end_labels_flat = end_labels.reshape(-1)
+
+        end_loss = self.ce(end_logits_flat, end_labels_flat)
+        end_valid = end_labels_flat != -100
+
+        if end_valid.any():
+            end_loss = end_loss[end_valid].mean()
+        else:
+            end_loss = logits.sum() * 0.0
+
+        return {
+            "total": loss + start_loss + end_loss,
+            "span_loss": loss,
+            "start_loss": start_loss,
+            "end_loss": end_loss,
+        }
+
+# %% [code]
+## Viết eval_fn vào đây
+
+# Bỏ hết eval_fn và trọng số vào đây
+class CustomEvalFn(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def compute_f1(self, tp, fp, fn):
+        precision = tp / (tp + fp + 1e-8)
+        recall = tp / (tp + fn + 1e-8)
+        f1 = 2 * precision * recall / (precision + recall + 1e-8)
+        return precision, recall, f1
+
+    def forward(self, pred, gold):
+        pred_set = set(pred)
+        gold_set = set(gold)
+
+        tp = len(pred_set & gold_set)
+        fp = len(pred_set - gold_set)
+        fn = len(gold_set - pred_set)
+
+        precision, recall, f1 = self.compute_f1(tp, fp, fn)
+
+        return {
+                f"precision": precision,
+                f"recall": recall,
+                f"f1": f1,
+            }
+
+class SpanErrorAnalyzer:
+    def __init__(self, pad_token_id=0):
+        self.pad_token_id = pad_token_id
+
+    # ===== helper =====
+    def _to_set(self, data):
+        """
+        data: list of (b, tuple(ids))
+        -> dict[b] = set(tuple(ids))
+        """
+        res = defaultdict(set)
+        for b, ids in data:
+            ids = tuple([i for i in ids if i != self.pad_token_id])
+            if len(ids) > 0:
+                res[b].add(ids)
+        return res
+
+    def _iou(self, a, b):
+        """
+        a, b: tuple(ids)
+        """
+        set_a, set_b = set(a), set(b)
+        inter = len(set_a & set_b)
+        union = len(set_a | set_b)
+        if union == 0:
+            return 0.0
+        return inter / union
+
+    def _boundary_error(self, pred, gold):
+        """
+        đo lệch boundary dựa trên overlap prefix/suffix
+        """
+        # left match
+        left = 0
+        for i in range(min(len(pred), len(gold))):
+            if pred[i] == gold[i]:
+                left += 1
+            else:
+                break
+
+        # right match
+        right = 0
+        for i in range(1, min(len(pred), len(gold)) + 1):
+            if pred[-i] == gold[-i]:
+                right += 1
+            else:
+                break
+
+        return {
+            "left_match": left,
+            "right_match": right,
+            "pred_len": len(pred),
+            "gold_len": len(gold),
+        }
+
+    # ===== main =====
+    def analyze(self, preds, golds):
+        pred_map = self._to_set(preds)
+        gold_map = self._to_set(golds)
+
+        all_batches = set(pred_map.keys()) | set(gold_map.keys())
+
+        stats = Counter()
+
+        detailed_errors = []
+
+        for b in all_batches:
+            pset = pred_map.get(b, set())
+            gset = gold_map.get(b, set())
+
+            matched_gold = set()
+
+            # ===== check predictions =====
+            for p in pset:
+                if p in gset:
+                    stats["exact_match"] += 1
+                    matched_gold.add(p)
+                else:
+                    # tìm gold gần nhất
+                    best_iou = 0
+                    best_g = None
+
+                    for g in gset:
+                        iou = self._iou(p, g)
+                        if iou > best_iou:
+                            best_iou = iou
+                            best_g = g
+
+                    if best_iou > 0:
+                        stats["partial_match"] += 1
+
+                        boundary = self._boundary_error(p, best_g)
+
+                        detailed_errors.append({
+                            "type": "boundary_error",
+                            "batch": b,
+                            "pred": p,
+                            "gold": best_g,
+                            "iou": best_iou,
+                            **boundary
+                        })
+                    else:
+                        if b not in gold_map:
+                            stats["no_event_sample"] += 1
+                            err_type = "no_event_sample"
+                        else:
+                            stats["completely_wrong"] += 1
+                            err_type = "completely_wrong"
+                    
+                        detailed_errors.append({
+                            "type": err_type,
+                            "batch": b,
+                            "pred": p
+                        })
+
+            # ===== check missing =====
+            for g in gset:
+                if g not in matched_gold:
+                    # check if any pred overlaps
+                    overlap = any(self._iou(p, g) > 0 for p in pset)
+
+                    if overlap:
+                        stats["miss_with_overlap"] += 1
+                    else:
+                        stats["miss"] += 1
+
+                    detailed_errors.append({
+                        "type": "miss",
+                        "batch": b,
+                        "gold": g
+                    })
+
+        return {
+            "summary": {
+                "exact_match": (stats["exact_match"], stats["exact_match"] / len(preds)),
+                "partial_match": (stats["partial_match"], stats["partial_match"] / len(preds)),
+                "no_event_sample": (stats["no_event_sample"], stats["no_event_sample"] / len(preds)),
+                "completely_wrong": (stats["completely_wrong"], stats["completely_wrong"] / len(preds)),
+                "miss": (stats["miss"], stats["miss"] / len(golds)),
+                "miss_with_overlap": (stats["miss_with_overlap"], stats["miss_with_overlap"] / len(golds)),
+            },
+            "details": detailed_errors
+        }
+
+# %% [code]
+class DataParallelProxy(nn.DataParallel):
+
+    def __getattr__(self, name):
+        try:
+            return super().__getattr__(name)
+
+        except AttributeError:
+
+            attr = getattr(self.module, name)
+
+            if callable(attr):
+
+                def wrapper(*args, **kwargs):
+                    return self._parallel_apply_method(
+                        name,
+                        *args,
+                        **kwargs
+                    )
+
+                return wrapper
+
+            return attr
+
+    # =========================================================
+    # parallel custom method
+    # =========================================================
+
+    def _parallel_apply_method(self, method_name, *inputs, **kwargs):
+
+        if not self.device_ids:
+            return getattr(self.module, method_name)(*inputs, **kwargs)
+
+        inputs_scattered, kwargs_scattered = self.scatter(
+            inputs,
+            kwargs,
+            self.device_ids
+        )
+
+        replicas = self.replicate(
+            self.module,
+            self.device_ids[:len(inputs_scattered)]
+        )
+
+        outputs = self.parallel_apply(
+            [getattr(replica, method_name) for replica in replicas],
+            inputs_scattered,
+            kwargs_scattered
+        )
+
+        return self._custom_gather(
+            outputs,
+            self.output_device
+        )
+
+    # =========================================================
+    # OVERRIDE FORWARD GATHER
+    # =========================================================
+
+    def gather(self, outputs, output_device):
+
+        return self._custom_gather(
+            outputs,
+            output_device
+        )
+
+    # =========================================================
+    # recursive gather
+    # =========================================================
+
+    def _custom_gather(self, outputs, output_device):
+
+        first = outputs[0]
+
+        # =====================================================
+        # tensor
+        # =====================================================
+
+        if torch.is_tensor(first):
+
+            return self._gather_tensor(
+                outputs,
+                output_device
+            )
+
+        # =====================================================
+        # tuple
+        # =====================================================
+
+        if isinstance(first, tuple):
+
+            return tuple(
+                self._custom_gather(
+                    list(items),
+                    output_device
+                )
+                for items in zip(*outputs)
+            )
+
+        # =====================================================
+        # list
+        # =====================================================
+
+        if isinstance(first, list):
+
+            # list[tensor]
+            if len(first) > 0 and torch.is_tensor(first[0]):
+
+                return self._gather_tensor_list(
+                    outputs,
+                    output_device
+                )
+
+            merged = []
+
+            for out in outputs:
+                merged.extend(out)
+
+            return merged
+
+        # =====================================================
+        # dict
+        # =====================================================
+
+        if isinstance(first, dict):
+
+            return {
+                k: self._custom_gather(
+                    [o[k] for o in outputs],
+                    output_device
+                )
+                for k in first.keys()
+            }
+
+        # =====================================================
+        # fallback
+        # =====================================================
+
+        return outputs
+
+    # =========================================================
+    # gather tensor with auto pad
+    # =========================================================
+
+    def _gather_tensor(self, tensors, output_device):
+
+        # move same device first
+        tensors = [
+            t.to(output_device)
+            for t in tensors
+        ]
+
+        # =====================================================
+        # fast path
+        # =====================================================
+
+        try:
+            return torch.cat(tensors, dim=0)
+
+        except RuntimeError:
+            pass
+
+        # =====================================================
+        # auto max shape
+        # =====================================================
+
+        max_shape = list(tensors[0].shape)
+
+        for t in tensors[1:]:
+
+            for d in range(len(max_shape)):
+
+                max_shape[d] = max(
+                    max_shape[d],
+                    t.shape[d]
+                )
+
+        # =====================================================
+        # pad tensors
+        # =====================================================
+
+        padded = []
+
+        for t in tensors:
+
+            pad = []
+
+            # reverse order for F.pad
+            for d in reversed(range(len(max_shape))):
+
+                # never pad batch dim
+                if d == 0:
+                    pad.extend([0, 0])
+                    continue
+
+                diff = max_shape[d] - t.shape[d]
+
+                pad.extend([0, diff])
+
+            t = F.pad(t, pad)
+
+            padded.append(t)
+
+        return torch.cat(padded, dim=0)
+
+    # =========================================================
+    # list[tensor]
+    # =========================================================
+
+    def _gather_tensor_list(self, outputs, output_device):
+
+        merged = []
+
+        for out in outputs:
+            merged.extend(out)
+
+        return self._gather_tensor(
+            merged,
+            output_device
+        )
+
+# %% [code]
+## Viết cấu trúc model vào đây
+def get_span_reprs(hidden, spans):
+    """
+    Args:
+        hidden: (B, L, H)
+        spans:  (B, N, 2)
+
+    Return:
+        span_repr: (B, N, 4*H)
+    """
+
+    B, N, _ = spans.shape
+    H = hidden.size(-1)
+
+    batch_idx = torch.arange(B, device=hidden.device).unsqueeze(1)
+
+    start_idx = spans[..., 0]  # (B, N)
+    end_idx   = spans[..., 1]  # (B, N)
+    start_h = hidden[batch_idx, start_idx]
+    end_h   = hidden[batch_idx, end_idx]
+
+    span_repr = torch.cat(
+        [start_h, end_h, end_h - start_h, end_h * start_h],
+        dim=-1
+    )
+
+    return span_repr
+
+def filter_spans(
+    start_logits,      # (B, L, C)
+    end_logits,        # (B, L, C)
+    attn_mask,         # (B, L)
+    max_span_length=10,
+    topk_spans=20,
+    keep_neighbor=1
+):
+    """
+    Return:
+        spans: (B, N, 2)
+            N là số span lớn nhất trong batch sau expand + remove duplicate.
+            Padding bằng (0, 0)
+    """
+
+    device = start_logits.device
+    B, L, C = start_logits.shape
+
+    # Clone
+    start_logits = start_logits.detach()
+    end_logits = end_logits.detach()
+
+    # Prob
+    prob_s = F.softmax(start_logits, dim=-1)  # (B, L, C)
+    prob_e = F.softmax(end_logits, dim=-1)    # (B, L, C)
+
+    all_batch_spans = []
+
+    for b in range(B):
+
+        valid_len = int(attn_mask[b].sum().item())
+
+        candidate_spans = []
+
+        # Enumerate all valid spans
+        for s in range(1, valid_len):
+
+            max_e = min(valid_len - 1, s + max_span_length - 1)
+
+            for e in range(s, max_e + 1):
+
+                length = e - s + 1
+
+                score = (
+                    (1.0 - prob_s[b, s, 0].item())
+                    + (1.0 - prob_e[b, e, 0].item())
+                    - (length / max_span_length)
+                )
+                candidate_spans.append((score, s, e))
+
+        # Top-k
+        candidate_spans = sorted(
+            candidate_spans,
+            key=lambda x: x[0],
+            reverse=True
+        )[:topk_spans]
+
+        # Expand neighbors
+        final_spans = set()
+
+        for _, s, e in candidate_spans:
+
+            for ds in range(-keep_neighbor, keep_neighbor + 1):
+                for de in range(-keep_neighbor, keep_neighbor + 1):
+
+                    ns = s + ds
+                    ne = e + de
+
+                    if ns <= 0:
+                        continue
+
+                    if ne >= valid_len:
+                        continue
+
+                    if ns > ne:
+                        continue
+
+                    if (ne - ns + 1) > max_span_length:
+                        continue
+
+                    final_spans.add((ns, ne))
+
+        final_spans = sorted(list(final_spans))
+
+        if len(final_spans) == 0:
+            final_spans = [(0, 0)]
+
+        all_batch_spans.append(final_spans)
+
+    # Padding
+    max_num_spans = max(len(x) for x in all_batch_spans)
+
+    padded_spans = []
+
+    for spans in all_batch_spans:
+
+        pad_size = max_num_spans - len(spans)
+
+        spans = spans + [(0, 0)] * pad_size
+
+        padded_spans.append(spans)
+
+    spans = torch.tensor(
+        padded_spans,
+        dtype=torch.long,
+        device=device
+    )  # (B, N, 2)
+
+    return spans
+
+class MLP(nn.Module):
+    def __init__(self, in_size, hid_size, out_size):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(in_size, hid_size),
+            nn.ReLU(),
+            nn.Linear(hid_size, out_size)
+        )
+
+    def forward(self, x):
+        return self.mlp(x)
+
+class IEModel(nn.Module):
+    def __init__(self, backbone_model_name, num_labels, max_span_len, topk_spans, keep_neighbor):
+        super().__init__()
+        self.max_span_len = max_span_len
+        self.topk_spans = topk_spans
+        self.keep_neighbor = keep_neighbor
+        
+        self.encoder = AutoModel.from_pretrained(backbone_model_name)
+        hidden_size = self.encoder.config.hidden_size
+
+        self.start_classifier = MLP(hidden_size, hidden_size, num_labels)
+        self.end_classifier = MLP(hidden_size, hidden_size, num_labels)
+
+        self.span_classifier = MLP(4*hidden_size, hidden_size, num_labels)
+
+    def encode(self, input_ids, attention_mask):
+        B, n_parts, L = input_ids.shape
+        input_ids = input_ids.view(-1, L)
+        attention_mask = attention_mask.view(-1, L)
+        
+        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
+        hidden_states = outputs.last_hidden_state # B * n_parts, L, H
+        
+        hidden_states = hidden_states.view(B, n_parts, L, -1).reshape(B, n_parts*L, -1) # B, L, H
+        attention_mask = attention_mask.view(B, n_parts, L).reshape(B, n_parts*L) # B, L, H
+        return hidden_states, attention_mask
+
+    def get_token_logits(self, hidden_states):
+        start_logits = self.start_classifier(hidden_states)      # B, N, classes
+        end_logits = self.end_classifier(hidden_states)      # B, N, classes
+        return start_logits, end_logits
+
+    def get_logits(self, span_reprs):
+        logits = self.span_classifier(span_reprs)      # N, classes
+        return logits
+
+    def forward(self, input_ids, attention_mask, spans=None):
+        hidden_states, attention_mask = self.encode(input_ids, attention_mask)
+        start_logits, end_logits = self.get_token_logits(hidden_states)
+        if spans is None:
+            spans = filter_spans(start_logits, end_logits, attention_mask, self.max_span_len, self.topk_spans, self.keep_neighbor)
+        span_reprs = get_span_reprs(hidden_states, spans)
+        logits = self.get_logits(span_reprs)
+        return start_logits, end_logits, logits, spans
+
+def test_model():
+    model = DataParallelProxy(IEModel(backbone_model_name, 7, 10, 100, 2)).to(device)
+    model.eval()
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"Total params: {total_params:,}")
+
+    vocab_size = model.module.encoder.config.vocab_size
+    max_len = model.module.encoder.config.max_position_embeddings
+
+    bz = 32
+    i = torch.randint(0, vocab_size, (bz, 5, 10)).to(device)
+    a = torch.ones(bz, 5, 10).to(device)
+    s = torch.ones(bz, 3, 2, dtype=torch.long).to(device)
+
+    with torch.no_grad():
+        r = model(i, a)
+
+    if type(r) == tuple:
+        print([r[i].shape if type(r[i]) == type(torch.Tensor()) else len(r[i]) for i in range(len(r))])
+    else:
+        print(r.shape)
+
+test_model()
+
+# %% [code]
+def configure_optimizers(network, optim_params, scheduler_params):
+    try:
+        optim_params = copy.copy(optim_params)
+        scheduler_params = copy.copy(scheduler_params)
+        
+        optim_name = optim_params.pop('name')
+        scheduler_name = scheduler_params.pop('name')
+    
+        optimizer_cls = globals().get(optim_name) or getattr(optim, optim_name, None)
+        scheduler_cls = globals().get(scheduler_name) or getattr(optim.lr_scheduler, scheduler_name, None)
+    
+        if optimizer_cls is None:
+            raise ValueError(f"Optimizer '{optim_name}' is not available!")
+    
+        optimizer = optimizer_cls(network.parameters(), **optim_params)
+        
+        scheduler = None
+        if scheduler_params and scheduler_cls:  # Chỉ tạo scheduler nếu có tham số
+            scheduler = scheduler_cls(optimizer, **scheduler_params)
+    
+        return optimizer, scheduler
+
+    except KeyError as e:
+        raise ValueError(f"Missing {e} in config!!")
+            
+def freeze(self, model):
+    model.eval()
+    for param in model.parameters():
+        param.requires_grad = False
+
+def unfreeze(self, model):
+    model.train()
+    for param in model.parameters():
+        param.requires_grad = True
+
+def reduce_batch_size(loader, ratio=0.5):
+    new_bs = max(1, int(loader.batch_size * ratio))
+
+    shuffle = isinstance(loader.sampler, RandomSampler)
+
+    new_loader = DataLoader(
+        dataset=loader.dataset,
+        batch_size=new_bs,
+        shuffle=shuffle,
+        sampler=None if shuffle else loader.sampler,
+        num_workers=loader.num_workers,
+        collate_fn=loader.collate_fn,
+        pin_memory=loader.pin_memory,
+        drop_last=loader.drop_last,
+        timeout=loader.timeout,
+        worker_init_fn=loader.worker_init_fn,
+        multiprocessing_context=loader.multiprocessing_context,
+        generator=loader.generator,
+        prefetch_factor=loader.prefetch_factor if loader.num_workers > 0 else None,
+        persistent_workers=loader.persistent_workers,
+        pin_memory_device=loader.pin_memory_device
+    )
+
+    return new_loader
+
+def list_to_tuple(x):
+    if isinstance(x, (list, tuple)):
+        return tuple(list_to_tuple(i) for i in x)
+    return x
+
+def fmt(x):
+    if isinstance(x, float):
+        return round(x, 5)
+    if isinstance(x, dict):
+        return {k: fmt(v) for k, v in x.items()}
+    if isinstance(x, list):
+        return [fmt(v) for v in x]
+    return x
+    
+class ModelEmaV3Proxy(ModelEmaV3):
+
+    def __getattr__(self, name):
+
+        try:
+            return super().__getattr__(name)
+
+        except AttributeError:
+
+            # tránh recursion
+            module = object.__getattribute__(self, "module")
+
+            return getattr(module, name)
+
+def align(
+    all_spans,   # (B, N, 2)
+    pred_spans,  # (B, M, 2)
+    obj,         # (B, N)
+    pad_value=0
+):
+    """
+    Return:
+        align_obj: (B, M)
+    """
+
+    B, N, _ = all_spans.shape
+    M = pred_spans.shape[1]
+
+    device = all_spans.device
+
+    # =========================================================
+    # Compare spans
+    # =========================================================
+
+    # (B, M, 1, 2)
+    pred_expand = pred_spans.unsqueeze(2)
+
+    # (B, 1, N, 2)
+    all_expand = all_spans.unsqueeze(1)
+
+    # (B, M, N)
+    match = (pred_expand == all_expand).all(dim=-1)
+
+    # =========================================================
+    # Gather obj
+    # =========================================================
+
+    # default = pad_value
+    align_obj = torch.full(
+        (B, M),
+        pad_value,
+        dtype=obj.dtype,
+        device=device
+    )
+
+    # matched index
+    has_match = match.any(dim=-1)  # (B, M)
+
+    matched_idx = match.float().argmax(dim=-1)  # (B, M)
+
+    gathered = torch.gather(
+        obj,
+        dim=1,
+        index=matched_idx
+    )  # (B, M)
+
+    align_obj[has_match] = gathered[has_match]
+
+    # =========================================================
+    # Force padding spans -> pad_value
+    # =========================================================
+
+    pad_mask = (pred_spans == 0).all(dim=-1)
+
+    align_obj[pad_mask] = pad_value
+
+    return align_obj
+
+def extract_spans(
+    all_spans,   # (B, N, 2)
+    all_label,   # (B, N)
+    pred_spans   # (B, M, 2)
+):
+    """
+    Return:
+        pred_list:
+            [(bidx, (s, e)), ...]
+
+        gold_list:
+            [(bidx, (s, e)), ...]
+    """
+
+    # =========================================================
+    # Gold spans
+    # =========================================================
+
+    gold_mask = all_label > 0  # (B, N)
+
+    gold_indices = gold_mask.nonzero(as_tuple=False)
+
+    gold_list = [
+        (
+            int(b),
+            (
+                int(all_spans[b, n, 0]),
+                int(all_spans[b, n, 1])
+            )
+        )
+        for b, n in gold_indices
+    ]
+
+    # =========================================================
+    # Pred spans
+    # =========================================================
+
+    pred_mask = (pred_spans > 0).all(dim=-1)  # (B, M)
+
+    pred_indices = pred_mask.nonzero(as_tuple=False)
+
+    pred_list = [
+        (
+            int(b),
+            (
+                int(pred_spans[b, m, 0]),
+                int(pred_spans[b, m, 1])
+            )
+        )
+        for b, m in pred_indices
+    ]
+
+    return gold_list, pred_list
+
+def extract_entities(
+    input_ids,    # (B, L)
+    start_logits, # (B, L, C)
+    end_logits,   # (B, L, C)
+    logits,       # (B, N, C)
+    pred_spans,   # (B, N, 2)
+    id2label,
+    alpha=1.0
+):
+    """
+    Return:
+    [
+        (batch_idx, ([token_ids], label_name)),
+        ...
+    ]
+    """
+
+    # =========================================================
+    # Log-softmax
+    # =========================================================
+
+    start_logprob = F.log_softmax(start_logits, dim=-1)  # (B, L, C)
+    end_logprob   = F.log_softmax(end_logits, dim=-1)    # (B, L, C)
+    span_logprob  = F.log_softmax(logits, dim=-1)        # (B, N, C)
+
+    # =========================================================
+    # Gather start/end score for pred spans
+    # =========================================================
+
+    start_idx = pred_spans[..., 0]  # (B, N)
+    end_idx   = pred_spans[..., 1]  # (B, N)
+
+    B, N = start_idx.shape
+    C = logits.shape[-1]
+
+    # (B, N, C)
+    start_score = torch.gather(
+        start_logprob,
+        dim=1,
+        index=start_idx.unsqueeze(-1).expand(-1, -1, C)
+    )
+
+    # (B, N, C)
+    end_score = torch.gather(
+        end_logprob,
+        dim=1,
+        index=end_idx.unsqueeze(-1).expand(-1, -1, C)
+    )
+
+    # =========================================================
+    # Ensemble score
+    # =========================================================
+
+    score = (
+        span_logprob
+        + alpha * (start_score + end_score)
+    )  # (B, N, C)
+
+    # =========================================================
+    # Predict label
+    # =========================================================
+
+    pred_labels = score.argmax(dim=-1)  # (B, N)
+
+    keep = (
+        (pred_labels > 0) &
+        (start_idx > 0) &
+        (end_idx > 0)
+    )
+
+    # =========================================================
+    # Extract entities
+    # =========================================================
+
+    results = []
+
+    for bidx in range(B):
+
+        valid_idxes = keep[bidx].nonzero(as_tuple=False).squeeze(-1)
+
+        for idx in valid_idxes:
+
+            lb = pred_labels[bidx, idx]
+
+            s, e = pred_spans[bidx, idx].tolist()
+
+            token_ids = input_ids[bidx, s:e+1].tolist()
+
+            results.append(
+                (
+                    bidx,
+                    (
+                        token_ids,
+                        id2label[lb.item()]
+                    )
+                )
+            )
+
+    return results
+    
+class Trainer:
+    def __init__(
+        self, training_time="00:11:30:00", eval_mode="max", topk=1, save_name="network", save_best=True, save_last=False, max_grad_norm=200.0, 
+        logging=0, logging_file=False, checkpoints_dir="", early_stopping=False, eval_from_ratio=-1, eval_every=1, device='cpu', 
+        schedule_in_step=True, use_ema=True, ema_from_ratio=-1, ema_decay=0.999, return_best=True, return_last=True
+    ):
+        self.ema_net = None
+        
+        self.training_time = self._time_str_to_seconds(training_time)
+        self.mode = eval_mode
+        self.topk = topk
+        self.device = device
+        self.logging = logging if logging < epochs else 1
+        self.logging_file = logging_file
+        self.checkpoints_dir = checkpoints_dir
+        self.early_stopping = early_stopping
+        self.eval_from_ratio = eval_from_ratio
+        self.eval_every = eval_every
+        self.save_name = save_name
+        self.save_best = save_best
+        self.save_last = save_last
+        self.return_best = return_best
+        self.return_last = return_last
+        self.max_grad_norm = max_grad_norm
+        self.schedule_in_step = schedule_in_step
+        self.use_ema = use_ema
+        self.ema_from_ratio = ema_from_ratio
+        self.ema_decay = ema_decay
+        
+        self.best_stage = [[float('-inf') if self.mode == 'max' else float('inf'), None, None]]
+        self.grad_scaler = torch.amp.GradScaler(self.device, init_scale=1024.0)
+
+    def fit(self, network, optimizer, scheduler, loss_fn, epochs, train_loader, val_loader=None, eval_fn=None, start_epoch=1, start_training_time=None, id2label=None):
+        if eval_fn is None:
+            if self.mode == "max":
+                eval_fn = lambda *x: -loss_fn(*x)
+            else:
+                eval_fn = lambda *x: loss_fn(*x)
+
+        if torch.cuda.device_count() > 1:
+            network = DataParallelProxy(network)
+        network = network.to(self.device)
+            
+        if not start_training_time:
+            start_training_time = time.time()
+            
+        start_ema = int(epochs * self.ema_from_ratio)
+        start_eval = int(epochs * self.eval_from_ratio)
+        
+        if val_loader is None:
+            print(f'[Trainer CallBack] 📢 Không có Val Set, không thể đánh giá và Early Stopping!')
+        else:
+            model_to_use_str = 'mô hình EMA' if self.use_ema else 'mô hình gốc'
+            start_model_update_str = f'Bắt đầu cập nhật EMA từ epoch {start_epoch + start_ema}!' if self.use_ema else ''
+            print(f'[Trainer CallBack] 📢 Đánh giá bằng {model_to_use_str} từ epoch {start_epoch + start_eval}!', start_model_update_str)
+
+        training_log = {}
+        for epoch in range(start_epoch, epochs+start_epoch):
+            if self.use_ema and self.ema_net is None and epoch - start_epoch >= start_ema:
+                self.ema_net = ModelEmaV3Proxy(network, self.ema_decay, device=self.device)
+                
+            try:
+                teaching_rate = math.cos(math.pi / 2 * epoch / epochs)
+                train_loss_epoch, train_loss_epoch_dict = self._train_epoch(network, train_loader, optimizer, scheduler, loss_fn, teaching_rate)
+                logging_dict = {'lr': [group['lr'] for group in optimizer.param_groups], 'train_loss': train_loss_epoch}
+                logging_dict.update(train_loss_epoch_dict)
+                    
+                if val_loader is not None and epoch - start_epoch >= start_eval and (epoch - start_epoch - start_eval) % self.eval_every == 0:
+                    eval_net = self.ema_net.module if (self.use_ema and self.ema_net is not None) else network
+                    
+                    val_score, val_score_dict, _ = self._eval_epoch(eval_net, val_loader, eval_fn, id2label)
+                    update = self._update_best_network(eval_net, val_score, epoch)
+                    logging_dict.update({'val_score': val_score, 'best_score': self.best_stage[0][0], 'new_best_model': update})
+                    logging_dict.update(val_score_dict)
+                if not self.schedule_in_step and scheduler:
+                    scheduler.step()
+                    
+            except RuntimeError as e:
+                if "out of memory" in str(e).lower():
+                    print(f"[Trainer CallBack] ⚠️ Epoch {epoch}/{epochs}: CUDA Out of Memory! Clearing GPU cache...")
+                    torch.cuda.empty_cache()
+                    gc.collect()
+                    if torch.cuda.is_available():
+                        torch.cuda.synchronize()
+                    print(f"[Trainer CallBack] ✅ Epoch {epoch}/{epochs}: GPU memory cleared")
+                    
+                    train_loader = reduce_batch_size(train_loader, ratio=0.5)
+                    if val_loader is not None:
+                        val_loader = reduce_batch_size(val_loader, ratio=0.5)
+                        
+                    logging_dict = {'lr': [group['lr'] for group in optimizer.param_groups], 'train_loss': float('inf')}
+                else:
+                    raise
+                    
+            training_log[epoch] = logging_dict                    
+            if self.is_early_stopping(epoch):
+                print(f'[Trainer CallBack] 📢 Epoch {epoch}/{epochs}: Detect Overfitting! Breaking Training Process...')
+                break
+            if self.logging:
+                if epoch % self.logging == 0:
+                    print(f'[Trainer CallBack] 📢 Epoch {epoch}/{epochs}:', fmt(logging_dict))
+            else:
+                print(f'{epoch}...', end=' ')
+                    
+            if self._at_time_limit(start_training_time):
+                print(f'[Trainer CallBack] ⚠️ Epoch {epoch}/{epochs}: Thời gian training giới hạn là {self.training_time}, hết giờ tại epoch {epoch}/{epochs}')
+                break
+                
+        if self.logging_file:
+            os.makedirs(f'{self.checkpoints_dir}/logs', exist_ok=True)
+            with open(f"{self.checkpoints_dir}/logs/{self.save_name}_logging.json", "a", encoding="utf-8") as f:
+                f.write(json.dumps(training_log))
+                
+        if self.use_ema and self.ema_net is not None:
+            self._save_state_dict(self.ema_net.module)
+        else:
+            self._save_state_dict(network)
+        print(f'[Trainer CallBack] 📢 Kết thúc training.\n')
+
+        best_model, last_model = None, None
+        eval_net = self.ema_net.module if (self.use_ema and self.ema_net is not None) else network
+        if self.return_best :
+            best_model = self.best_stage[0][2] if self.best_stage[0][2] is not None else eval_net.state_dict()
+            best_model = {k.replace("module.", ""): v.detach().cpu().clone() for k, v in best_model.items()}
+        if self.return_last:
+            last_model = eval_net.state_dict()
+            last_model = {k.replace("module.", ""): v.detach().cpu().clone() for k, v in last_model.items()}
+
+        del network
+        torch.cuda.empty_cache()
+        gc.collect()
+        return training_log, best_model, last_model
+
+    def _time_str_to_seconds(self, time_str):
+        days, hours, minutes, seconds = map(int, time_str.split(":"))
+        return days * 86400 + hours * 3600 + minutes * 60 + seconds
+
+    def _update_best_network(self, network, val_score, epoch):
+        topk = max(1, self.topk)
+        self.best_stage.append([val_score, epoch, {k: v.detach().cpu().clone() for k, v in network.state_dict().items()}])
+        self.best_stage = sorted(self.best_stage, reverse=(self.mode == 'max'), key=lambda x: x[0])[:topk]
+        if val_score in [x[0] for x in self.best_stage]:
+            return True
+        return False
+
+    def is_early_stopping(self, epoch):
+        if self.best_stage[0][1] is None:
+            return False
+        if not self.early_stopping:
+            return False
+        return epoch - self.best_stage[0][1] >= self.early_stopping
+
+    def _at_time_limit(self, start_training_time):
+        return time.time() - start_training_time >= self.training_time
+
+    def _save_state_dict(self, network):
+        if self.topk <= 0:
+            return
+
+        if self.save_best:
+            for r in range(self.topk):
+                os.makedirs(f'{self.checkpoints_dir}/r{r+1}s', exist_ok=True)
+                
+            for rank, (score, epoch, state_dict) in enumerate(self.best_stage):
+                if state_dict is None:
+                    continue
+                state_dict = {k.replace("module.", ""): v.detach().cpu().clone() for k, v in state_dict.items()}
+                torch.save(state_dict, f'{self.checkpoints_dir}/r{rank+1}s/{self.save_name}_r{rank+1}_vs{score:.5f}_{"ema" if self.ema_net is not None else ""}.pth')
+        if self.save_last:
+            os.makedirs(f'{self.checkpoints_dir}/lasts', exist_ok=True)
+            state_dict = {k.replace("module.", ""): v.detach().cpu().clone() for k, v in network.state_dict().items()}
+            torch.save(state_dict, f'{self.checkpoints_dir}/lasts/{self.save_name}_last_{"ema" if self.ema_net is not None else ""}.pth')
+        
+    def _train_epoch(self, network, train_loader, optimizer, scheduler, loss_fn, teaching_rate):
+        network.train()
+        total_loss = 0
+        total_loss_dict = {}
+        for batch_idx, batch in enumerate(train_loader):
+            optimizer.zero_grad()
+            with torch.autocast(device_type=self.device, dtype=torch.float16):
+                loss, loss_dict = self._cal_loss(network, batch, batch_idx, loss_fn, teaching_rate)
+
+                for k, v in loss_dict.items():
+                    t = total_loss_dict.get(k, 0)
+                    total_loss_dict[k] = t + v
+            self.grad_scaler.scale(loss).backward()
+            self.grad_scaler.unscale_(optimizer)
+            grad_norm = nn.utils.clip_grad_norm_(network.parameters(), self.max_grad_norm)
+            # print(grad_norm) # Bỏ cmt dòng này để biết nên chọn max_grad_norm bằng bao nhiêu...
+            self.grad_scaler.step(optimizer)
+            self.grad_scaler.update()
+            if self.schedule_in_step and scheduler:
+                scheduler.step()
+            if self.use_ema and self.ema_net is not None:
+                self.ema_net.update(network)
+            total_loss += loss
+        return (total_loss / len(train_loader)).item(), {k: v.item() / len(train_loader) for k, v in total_loss_dict.items()}
+
+    def _eval_epoch(self, network, val_loader, eval_fn, id2label):
+        network.eval()
+        total_score = 0.0
+        total_score_dict = {}
+        object_lists = None  # sẽ init sau
+    
+        with torch.no_grad():
+            for batch_idx, batch in enumerate(val_loader):
+                score, score_dict, objects = self._cal_val_score(network, batch, batch_idx, eval_fn, id2label)
+                total_score += score
+
+                for k, v in score_dict.items():
+                    t = total_score_dict.get(k, 0)
+                    total_score_dict[k] = t + v
+    
+                if objects:
+                    if object_lists is None:
+                        object_lists = [[] for _ in range(len(objects))]
+    
+                    for i, obj in enumerate(objects):
+                        object_lists[i].append(obj.detach())
+    
+        if object_lists is not None:
+            object_arrays = [
+                torch.concat(obj_list, dim=0).cpu().numpy()
+                for obj_list in object_lists
+            ]
+        else:
+            object_arrays = []
+    
+        return total_score / len(val_loader), {k: v / len(val_loader) for k, v in total_score_dict.items()}, object_arrays
+
+    def _cal_loss(self, network, batch, batch_idx, loss_fn, teaching_rate):
+        # Bạn cần override _cal_loss để tính loss
+        input_ids = batch['input_ids'].to(self.device)
+        attention_mask = batch['attention_mask'].to(self.device)
+        all_spans = batch['all_spans'].to(self.device)
+        all_labels = batch['all_labels'].to(self.device)
+        all_weights = batch['all_weights'].to(self.device)
+        start_labels = batch['start_labels'].to(self.device)
+        end_labels = batch['end_labels'].to(self.device)
+
+        choice = random.random()
+        if choice < teaching_rate:
+            start_logits, end_logits, logits, pred_spans = network(input_ids, attention_mask, all_spans)
+        else:
+            start_logits, end_logits, logits, pred_spans = network(input_ids, attention_mask)
+
+        align_labels = align(all_spans, pred_spans, all_labels, -100)
+        align_weights = align(all_spans, pred_spans, all_weights, 0)
+        
+        loss_dict = loss_fn(
+            logits, align_labels, align_weights, 
+            start_logits, start_labels, 
+            end_logits, end_labels,
+        )
+        return loss_dict['total'], loss_dict
+
+    def _cal_val_score(self, network, batch, batch_idx, eval_fn, id2label):
+        # Bạn cần override _cal_val_score để tính val score, list bên cạnh là để trả về y hay pred gì đó (nếu cần)
+        input_ids = batch['input_ids'].to(self.device)
+        attention_mask = batch['attention_mask'].to(self.device)
+        all_spans = batch['all_spans'].to(self.device)
+        all_labels = batch['all_labels'].to(self.device)
+        gold_entities = batch['gold_entities']
+
+        B, _, _ = input_ids.shape
+
+        start_logits, end_logits, logits, pred_spans = network(input_ids, attention_mask)
+        
+        gold_list, pred_list = extract_spans(all_spans, all_labels, pred_spans)
+        gold_list = list_to_tuple(gold_list)
+        pred_list = list_to_tuple(pred_list)
+        span_score = eval_fn(pred_list, gold_list)['recall']
+        
+        pred_ids = extract_entities(input_ids.reshape(B, -1), start_logits, end_logits, logits, pred_spans, id2label, alpha=0.5)
+        pred_ids = list_to_tuple(pred_ids)            
+        gold_ids = list_to_tuple(gold_entities)        
+        score_dict = eval_fn(pred_ids, gold_ids)
+
+        score_dict.update({'span_recall': span_score})
+        return score_dict['f1'] + score_dict['span_recall'], score_dict, []
+
+# %% [code]
+class PhoBERTSpanAligner:
+    def __init__(self, tokenizer, max_len):
+        self.tokenizer = tokenizer
+        self.max_len = max_len
+
+    # ===== 1. Extract discontinuous spans =====
+    def extract_spans(self, sample):
+        entity_spans = []
+        
+        for event in sample["entities"]:
+            entity_type = event["label"]
+            spans = [tuple(event["offset"])]
+            entity_spans.append({
+                "spans": spans,
+                "label": entity_type
+            })
+        
+        return entity_spans
+
+    # ===== 2. Word offsets =====
+    def build_word_offsets(self, text, words):
+        offsets = []
+        pointer = 0
+
+        for word in words:
+            start = text.find(word, pointer)
+            end = start + len(word)
+            offsets.append((start, end))
+            pointer = end
+
+        return offsets
+
+    # ===== 3. Char → word =====
+    def char_span_to_word_span(self, word_offsets, start, end):
+        start_word = None
+        end_word = None
+
+        for i, (w_start, w_end) in enumerate(word_offsets):
+            if w_start <= start < w_end:
+                start_word = i
+            if w_start < end <= w_end:
+                end_word = i
+
+        return start_word, end_word
+
+    # ===== 4. Word → subword =====
+    def word_to_subword_map(self, words):
+        mapping = []
+        subword_index = 1  # <s>
+
+        for word in words:
+            sub_tokens = self.tokenizer.tokenize(word)
+            start = subword_index
+            end = subword_index + len(sub_tokens) - 1
+            mapping.append((start, end))
+            subword_index += len(sub_tokens)
+
+        return mapping
+
+    # ===== 5. Span → subword =====
+    def span_to_subword(self, word_offsets, word_subword_map, spans):
+        sub_spans = []
+
+        for span_start, span_end in spans:
+            w_start, w_end = self.char_span_to_word_span(
+                word_offsets, span_start, span_end
+            )
+            if w_start is None or w_end is None:
+                continue
+
+            sub_start = word_subword_map[w_start][0]
+            sub_end = word_subword_map[w_end][1]
+            sub_spans.append((sub_start, sub_end))
+
+        return sub_spans
+
+    def extract_valid_spans(self, sub_spans):
+        valid_spans = []
+        for s, e in sub_spans:
+            if s < 0 or e < 0 or s >= self.max_len or e >= self.max_len or s > e:
+                continue
+            valid_spans.append((s, e))
+        return valid_spans
+
+    def encode(self, sample):
+        text = sample["text"]
+        entities = self.extract_spans(sample)
+    
+        # ===== 1. Word tokenize =====
+        words = word_tokenize(text)
+        sentence = " ".join(words)
+    
+        # ===== 2. Mapping =====
+        word_offsets = self.build_word_offsets(text, words)
+        word_subword_map = self.word_to_subword_map(words)
+    
+        # ===== 3. Tokenize FULL =====
+        encoding = self.tokenizer(
+            sentence,
+            max_length=self.max_len,
+            truncation=True,
+            padding="max_length",
+            return_tensors="pt"
+        )
+        input_ids = encoding["input_ids"][0]
+        attention_mask = encoding["attention_mask"][0]
+    
+        # ===== 5. Convert spans =====
+        entities_gold_spans = []
+    
+        for ent in entities:
+            label = ent["label"]
+            
+            sub_spans = self.span_to_subword(
+                word_offsets,
+                word_subword_map,
+                ent["spans"]
+            )
+            valid_spans = self.extract_valid_spans(sub_spans)
+            if len(valid_spans) == 0:
+                continue    
+            entities_gold_spans.append((tuple(valid_spans), label))
+                
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "entities_gold_spans": entities_gold_spans,
+        }
+
+def generate_spans(attention_mask, max_span_len):
+    seq_len = attention_mask.sum().item() - 2
+    spans = []
+    for i in range(1, seq_len+1):
+        for j in range(i, min(i+max_span_len, seq_len+1)):
+            spans.append((i, j))
+    return spans
+
+def match_gold_labels(
+    gold_spans,      # (N, 2)
+    gold_labels,     # (N,)
+    pred_spans,      # (M, 2)
+    default_label=-100
+):
+    """
+    Return:
+        pred_labels: (M,)
+    """
+
+    pred_labels = torch.full(
+        (pred_spans.size(0),),
+        default_label,
+        dtype=gold_labels.dtype,
+        device=gold_labels.device
+    )
+    if gold_spans.size(0) == 0:
+        return pred_labels
+
+    # (M, N)
+    matched = (pred_spans[:, None, :] == gold_spans[None, :, :]).all(dim=-1)
+    has_match = matched.any(dim=1)
+
+    # lấy index gold đầu tiên match
+    gold_idx = matched.float().argmax(dim=1)
+
+    pred_labels[has_match] = gold_labels[gold_idx[has_match]]
+
+    return pred_labels
+    
+class KLTNDataset(Dataset):
+    def __init__(self, all_data, using_idxes, label2id, tokenizer, max_len, max_n_parts, max_span_len, weight_rate):
+        super().__init__()
+        self.tokenizer = tokenizer
+        self.aligner = PhoBERTSpanAligner(tokenizer, max_len*max_n_parts)
+        self.all_data = all_data
+        self.using_idxes = using_idxes
+        self.label2id = label2id
+        self.max_len = max_len
+        self.max_n_parts = max_n_parts
+        self.max_span_len = max_span_len
+        self.weight_rate = weight_rate
+
+    def __len__(self):
+        return len(self.using_idxes)
+
+    def span_iou(self, span1, span2):
+        s1, e1 = span1
+        s2, e2 = span2
+    
+        # intersection
+        inter_left = max(s1, s2)
+        inter_right = min(e1, e2)
+        intersection = max(0, inter_right - inter_left + 1)
+    
+        # lengths
+        len1 = e1 - s1 + 1
+        len2 = e2 - s2 + 1
+    
+        # union
+        union = len1 + len2 - intersection
+        if union == 0:
+            return 0.0
+    
+        return intersection / union
+
+    def get_weights(self, spans, pos_spans):
+        # spans: (N, 2), pos_spans: (K, 2)
+        N, K = spans.size(0), pos_spans.size(0)
+        device = spans.device
+    
+        # ===== edge case =====
+        if K == 0:
+            weights = torch.ones(N, device=device, dtype=torch.float)
+            return weights
+    
+        # ===== IoU =====
+        s1, e1 = spans[:, None, 0], spans[:, None, 1]
+        s2, e2 = pos_spans[None, :, 0], pos_spans[None, :, 1]
+    
+        inter_s = torch.maximum(s1, s2)
+        inter_e = torch.minimum(e1, e2)
+        inter = (inter_e - inter_s + 1).clamp(min=0)
+    
+        len1 = (e1 - s1 + 1)
+        len2 = (e2 - s2 + 1)
+        union = len1 + len2 - inter
+    
+        iou = inter / (union + 1e-8)   # (N, K)
+    
+        # ===== weights: IoU=0 -> 1, else 10*IoU =====
+        max_iou, _ = iou.max(dim=1)
+        if self.weight_rate is not None:
+            weights = torch.where(max_iou > 0, 1 + self.weight_rate * max_iou, torch.ones_like(max_iou))
+        else:
+            weights = torch.ones_like(max_iou)
+    
+        return weights
+        
+    def to_span_label_tensors(self, data, label_map):
+        if len(data) == 0:
+            return (
+                torch.zeros((0, 2), dtype=torch.long),
+                torch.zeros((0,), dtype=torch.long)
+            )
+        
+        spans = torch.tensor([list(spans[0]) for spans, _ in data], dtype=torch.long)
+        labels = torch.tensor([label_map[label] for _, label in data], dtype=torch.long)
+        return spans, labels
+
+    def __getitem__(self, idx):
+        ridx = self.using_idxes[idx]
+        sample = self.all_data[ridx]
+        result = self.aligner.encode(sample)
+        
+        input_ids = result["input_ids"].squeeze(0)
+        attention_mask = result["attention_mask"].squeeze(0)
+        entities_gold_spans = result["entities_gold_spans"]
+
+        # Get all spans
+        all_spans = torch.tensor(generate_spans(attention_mask, self.max_span_len))
+        gold_spans = torch.tensor([spans[0] for spans, _ in entities_gold_spans], dtype=torch.long) if entities_gold_spans else torch.empty(0, 2, dtype=torch.long)
+        gold_labels = torch.tensor([self.label2id[label] for _, label in entities_gold_spans], dtype=torch.long) if entities_gold_spans else torch.empty(0, dtype=torch.long)
+        all_labels = match_gold_labels(
+            gold_spans,      # (N, 2)
+            gold_labels,     # (N,)
+            all_spans,      # (M, 2)
+            default_label=0
+        )
+        all_weights = self.get_weights(all_spans, gold_spans)
+
+        # Get label
+        gold_entities = []
+        start_labels = torch.ones_like(input_ids) * (1-attention_mask) * (-100)
+        end_labels = torch.ones_like(input_ids) * (1-attention_mask) * (-100)
+        for spans, label in entities_gold_spans:
+            s, e = spans[0]
+            
+            start_labels[s] = self.label2id[f'{label}']
+            end_labels[e] = self.label2id[f'{label}']
+            
+            gold_entities.append((tuple(input_ids[s:e+1].tolist()), label))
+        
+        input_ids = input_ids.reshape(self.max_n_parts, self.max_len)
+        attention_mask = attention_mask.reshape(self.max_n_parts, self.max_len)
+
+        n_valid_parts = math.ceil(attention_mask.sum().item() / self.max_len)
+        input_ids = input_ids[:n_valid_parts]
+        attention_mask = attention_mask[:n_valid_parts]
+        start_labels = start_labels[:n_valid_parts*self.max_len]
+        end_labels = end_labels[:n_valid_parts*self.max_len]
+    
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "all_spans": all_spans,
+            "all_labels": all_labels,
+            "all_weights": all_weights,
+            "start_labels": start_labels, 
+            "end_labels": end_labels, 
+            "gold_entities": gold_entities,
+        }
+
+def _pad_batch(tensor_list, pad_value=0):
+    """
+    tensor_list: list of tensors
+        mỗi tensor shape: (Nk, n_parts_i, max_len_i)
+
+    return:
+        padded tensor shape: (B, max_Nk, max_n_parts, max_len)
+    """
+
+    # lấy max toàn batch
+    max_Nk = max(t.size(0) for t in tensor_list)
+    max_n_parts = max(t.size(1) for t in tensor_list)
+    max_len = max(t.size(2) for t in tensor_list)
+
+    padded = []
+
+    for t in tensor_list:
+        Nk, n_parts_i, max_len_i = t.shape
+
+        # pad chiều n_parts và max_len trước
+        if n_parts_i < max_n_parts or max_len_i < max_len:
+            new_t = t.new_full(
+                (Nk, max_n_parts, max_len),
+                pad_value
+            )
+            new_t[:, :n_parts_i, :max_len_i] = t
+            t = new_t
+
+        # pad chiều Nk
+        if Nk < max_Nk:
+            pad_tensor = t.new_full(
+                (max_Nk - Nk, max_n_parts, max_len),
+                pad_value
+            )
+            t = torch.cat([t, pad_tensor], dim=0)
+
+        padded.append(t)
+
+    return torch.stack(padded)   # (B, max_Nk, max_n_parts, max_len)
+
+def collate_fn(batch):    
+    gold_entities = []
+    for bidx, b in enumerate(batch):
+        for entity in b['gold_entities']:
+            gold_entities.append([bidx, entity])
+
+    input_ids = [b["input_ids"].unsqueeze(-1) for b in batch]
+    attention_mask = [b["attention_mask"].unsqueeze(-1) for b in batch]
+    all_spans = [b["all_spans"].unsqueeze(-1) for b in batch]
+    all_labels = [b["all_labels"].unsqueeze(-1).unsqueeze(-1) for b in batch]
+    all_weights = [b["all_weights"].unsqueeze(-1).unsqueeze(-1) for b in batch]
+    start_labels = [b["start_labels"].unsqueeze(-1).unsqueeze(-1) for b in batch]
+    end_labels = [b["end_labels"].unsqueeze(-1).unsqueeze(-1) for b in batch]
+
+    # pad theo Nk
+    input_ids = _pad_batch(input_ids, pad_value=0).squeeze(-1)
+    attention_mask = _pad_batch(attention_mask, pad_value=0).squeeze(-1)
+    all_spans = _pad_batch(all_spans, pad_value=0).squeeze(-1)
+    all_labels = _pad_batch(all_labels, pad_value=-100).squeeze(-1).squeeze(-1)
+    all_weights = _pad_batch(all_weights, pad_value=0).squeeze(-1).squeeze(-1)
+    start_labels = _pad_batch(start_labels, pad_value=-100).squeeze(-1).squeeze(-1)
+    end_labels = _pad_batch(end_labels, pad_value=-100).squeeze(-1).squeeze(-1)
+    
+    return {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        "all_spans": all_spans,
+        "all_labels": all_labels,
+        "all_weights": all_weights,
+        "start_labels": start_labels,
+        "end_labels": end_labels,
+        "gold_entities": gold_entities,
+    }
+
+# %% [code]
+def shift_bidx(spans, batch_idx):
+    shifted = []
+    for bidx, ent in spans:
+        new_bidx = bidx + batch_idx * batch_size
+        shifted.append((new_bidx, ent))
+    return shifted
+
+def refactor_entities(entities, save_dict):
+    i, c = [], []
+    for bidx, (ids, lb) in entities:
+        if (bidx, ids) not in i:
+            i.append((bidx, ids))
+            
+        if (bidx, (ids, lb)) not in c:
+            c.append((bidx, (ids, lb)))
+
+    save_dict['Ent-I'].extend(i)
+    save_dict['Ent-C'].extend(c)
+
+def test(
+    network, state_dicts, test_loader, eval_fn, analyzer, device, id2label, tokenizer, 
+    alphas=[0.0, 0.2, 0.4, 0.5, 0.6, 0.8, 1.0], keep_neighbors=[0, 1, 2, 3]
+):
+    if torch.cuda.device_count() > 1:
+        network = DataParallelProxy(network)
+    network = network.to(device)
+    network.eval()
+
+    eval_types = ['Ent-I', 'Ent-C']
+    
+    all_pred = {(keep_neighbor, alpha): {eval_type: [] for eval_type in eval_types} for alpha in alphas for keep_neighbor in keep_neighbors}
+    all_gold = {eval_type: [] for eval_type in eval_types}
+
+    list_input_ids = []
+
+    with torch.no_grad():
+        for batch_idx, batch in enumerate(test_loader):
+            input_ids = batch['input_ids'].to(device)
+            attention_mask = batch['attention_mask'].to(device)
+            all_spans = batch['all_spans'].to(device)
+            gold_entities = batch['gold_entities']
+
+            B, _, _ = input_ids.shape
+            list_input_ids.extend(input_ids.reshape(B, -1).tolist())
+
+            list_hidden_states = []
+            list_start_logits = []
+            list_end_logits = []
+            for sd in state_dicts:
+                if torch.cuda.device_count() > 1:
+                    network.module.load_state_dict(sd)
+                else:
+                    network.load_state_dict(sd)
+                
+                hidden_states, attention_mask = network.encode(input_ids, attention_mask)
+                start_logits, end_logits = network.get_token_logits(hidden_states)
+                list_hidden_states.append(hidden_states)
+                list_start_logits.append(start_logits)
+                list_end_logits.append(end_logits)
+                
+            ensemble_start_logits = torch.stack(list_start_logits, dim=0).mean(dim=0)
+            ensemble_end_logits = torch.stack(list_end_logits, dim=0).mean(dim=0)
+
+            for keep_neighbor in keep_neighbors:
+                list_logits = []
+                spans = filter_spans(ensemble_start_logits, ensemble_end_logits, attention_mask, network.max_span_len, network.topk_spans, keep_neighbor)
+                
+                for sd, hidden_states in zip(state_dicts, list_hidden_states):
+                    if torch.cuda.device_count() > 1:
+                        network.module.load_state_dict(sd)
+                    else:
+                        network.load_state_dict(sd)
+                    span_reprs = get_span_reprs(hidden_states, spans)
+                    logits = network.get_logits(span_reprs)
+                    list_logits.append(logits)
+
+                ensemble_logits = torch.stack(list_logits, dim=0).mean(dim=0)
+                for alpha in alphas:
+                    pred_entities = extract_entities(input_ids.reshape(B, -1), ensemble_start_logits, ensemble_end_logits, ensemble_logits, spans, id2label, alpha=alpha)
+                    pred_entities = shift_bidx(pred_entities, batch_idx)
+                    refactor_entities(pred_entities, all_pred[keep_neighbor, alpha])
+                    
+            gold_entities = shift_bidx(gold_entities, batch_idx)
+            refactor_entities(gold_entities, all_gold)
+
+    # ===== GLOBAL EVAL =====
+    final_score = {}
+    max_key = -1
+    max_scores = -1
+    for key in all_pred.keys():
+        final_score[key] = {}
+        for eval_type in eval_types:
+            score = eval_fn(list_to_tuple(all_pred[key][eval_type]), list_to_tuple(all_gold[eval_type]))
+            final_score[key][eval_type] = score
+        if max_scores < final_score[key]['Ent-C']['f1']:
+            max_scores = final_score[key]['Ent-C']['f1']
+            max_key = key
+            
+    analyze_result = analyzer.analyze(list_to_tuple(all_pred[max_key]['Ent-I']),  list_to_tuple(all_gold['Ent-I']))
+
+    # ===== PREDICT =====
+    predictions = []
+    for input_ids in list_input_ids:
+        predictions.append([tokenizer.decode(input_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)])
+    for bidx, (ids, lb) in all_pred[max_key]['Ent-C']:
+        predictions[bidx].append((tokenizer.decode(ids, skip_special_tokens=True, clean_up_tokenization_spaces=True), lb))
+
+    return final_score, analyze_result, predictions
+
+# %% [code]
+with open(f'{train_dir}/train.json', "r", encoding="utf-8") as f:
+    data_train = json.load(f)
+    
+with open(f'{test_dir}/test.json', "r", encoding="utf-8") as f:
+    data_test = json.load(f)
+
+print('Train:', len(data_train))
+print('Test:', len(data_test))
+
+# %% [code]
+entity_types = ['O'] + sorted(list(set([e['label'] for d in data_train + data_test for e in d['entities']])))
+# bio_entity_type = ['O'] + [f'{prefix}-{ent}' for ent in entity_types for prefix in ['B', 'I']]
+label2id = {l: i for i, l in enumerate(entity_types)}
+id2label = {i: l for l, i in label2id.items()}
+
+# %% [code]
+zero_entities_idxes = []
+for idx, d in enumerate(data_train):
+    if len(d['entities']) == 0:
+        zero_entities_idxes.append(idx)
+
+n_zero_entities_samples = len(zero_entities_idxes)
+n_has_entities_samples = len(data_train) - n_zero_entities_samples
+
+random.seed(42)
+k = min(int(n_has_entities_samples * zero_entities_rate), len(zero_entities_idxes))
+sampled_zero_entities_idxes = random.sample(zero_entities_idxes, k)
+
+new_data_train = []
+for idx, d in enumerate(data_train):
+    if len(d['entities']) == 0:
+        if idx in sampled_zero_entities_idxes:
+            new_data_train.append(d)
+    else:
+        new_data_train.append(d)
+data_train = new_data_train
+
+print('Train:', len(data_train))
+
+# %% [code]
+if debug_only:
+    data_train = data_train[:10]
+    data_test = data_test[:10]
+
+    print('Train:', len(data_train))
+    print('Test:', len(data_test))
+
+# %% [code]
+tokenizer = AutoTokenizer.from_pretrained(backbone_model_name)
+
+# %% [code]
+print('Experiment name:', state_dict_save_name)
+
+# %% [code]
+# trainset = KLTNDataset(data_train, np.array(range(len(data_train))), label2id, tokenizer, **train_memory_params)
+# train_loader = DataLoader(trainset, collate_fn=collate_fn, **train_loader_params)
+# for b in train_loader:
+#     break
+
+# %% [code]
+if not test_only:
+    full_idxes = np.array(range(len(data_train)))
+    training_logs, best_models, last_models = [], [], []
+    start_training_time = time.time()
+    for seed in SEEDS:
+        kf = KFold(n_splits=nfolds, shuffle=True, random_state=seed)
+        for fold_idx, (tr_idx, va_idx) in enumerate(kf.split(full_idxes)):
+            if only_fold_idx is not None and only_fold_idx >= 0 and only_fold_idx != fold_idx:
+                continue
+            set_seed(seed)
+    
+            train_idxes, val_idxes = full_idxes[tr_idx], full_idxes[va_idx]
+    
+            trainset = KLTNDataset(data_train, train_idxes, label2id, tokenizer, **train_memory_params)
+            valset = KLTNDataset(data_train, val_idxes, label2id, tokenizer, **val_memory_params)
+            
+            generator = torch.Generator()
+            generator.manual_seed(seed)
+            train_loader = DataLoader(trainset, generator=generator, collate_fn=collate_fn, **train_loader_params)
+            val_loader = DataLoader(valset, generator=generator, collate_fn=collate_fn, **val_loader_params)
+    
+            my_model = IEModel(
+                num_labels=len(label2id), 
+                **model_params
+            )
+            total_params = sum(p.numel() for p in my_model.parameters())
+            print(f"Total params: {total_params:,}")
+            
+            # optimizer, scheduler = configure_optimizers(my_model, optim_params, scheduler_params)
+            encoder_params = set(map(id, my_model.encoder.parameters()))
+            other_params = [
+                p for p in my_model.parameters()
+                if id(p) not in encoder_params
+            ]        
+            optimizer = optim.AdamW([
+                {"params": my_model.encoder.parameters(), "lr": 2e-5},
+                {"params": other_params}
+            ], lr=5e-4)
+            scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=20, eta_min=1e-6)
+            
+            loss_fn = CustomLoss(
+                **loss_func_params
+            )
+            eval_fn = CustomEvalFn(**eval_func_params)
+            trainer_params['save_name'] = f'{state_dict_save_name}_s{seed}_f{fold_idx}'
+            trainer = Trainer(**trainer_params)
+    
+            print(f'Start Training Fold {fold_idx}...')
+            training_log, best_model, last_model = trainer.fit(
+                my_model, optimizer, scheduler, loss_fn, epochs, train_loader, val_loader, eval_fn, 
+                start_epoch=1, start_training_time=start_training_time, id2label=id2label
+            )
+            
+            training_logs.append(training_log)
+            best_models.append(best_model)
+            last_models.append(last_model)
+
+# %% [code]
+def load_all_state_dicts(folder):
+    files = []
+    
+    for file in os.listdir(folder):
+        if file.endswith(".pt") or file.endswith(".pth"):
+            m = re.search(r"f(\d+)", file)   # tìm f<số>
+            if m:
+                fold = int(m.group(1))
+                files.append((fold, file))
+
+    # sort theo fold
+    files.sort(key=lambda x: x[0])
+
+    state_dicts = []
+    for fold, file in files:
+        path = os.path.join(folder, file)
+        print(f"Loading fold {fold}: {file}")
+        state_dict = torch.load(path, map_location="cpu")
+        state_dicts.append(state_dict)
+
+    return state_dicts
+
+if test_only:
+    snapshot_download(repo_id=repo_name, local_dir="", repo_type="model", allow_patterns=[f"{state_dict_save_name}/**"])
+    get_ipython().system('rm -rf .cache .gitattributes')
+    
+    best_models = load_all_state_dicts(f"{state_dict_save_name}/r1s")
+    last_models = load_all_state_dicts(f"{state_dict_save_name}/lasts")
+
+# %% [code]
+def dict_to_df(data, row_names):
+    """
+    Input:
+    {
+        model_name: {
+            (level1, ..., level(n-1)): {
+                leveln: {
+                    metric: value
+                }
+            }
+        }
+    }
+
+    Output:
+    - level1 -> level(n-1): columns thường
+    - MultiIndex columns:
+        (model_name, leveln, metric)
+    """
+
+    rows = {}
+
+    for model_name, model_data in data.items():
+
+        for upper_levels, last_level_dict in model_data.items():
+
+            # đảm bảo tuple
+            if not isinstance(upper_levels, tuple):
+                upper_levels = (upper_levels,)
+
+            # ===== tạo key row =====
+            row_key = upper_levels
+
+            if row_key not in rows:
+
+                row = {}
+
+                for i, lv in enumerate(upper_levels):
+                    row[row_names[i]] = lv
+
+                rows[row_key] = row
+
+            # ===== add metrics =====
+            for last_level, metrics in last_level_dict.items():
+
+                for metric, value in metrics.items():
+
+                    rows[row_key][
+                        (model_name, last_level, metric)
+                    ] = value
+
+    # ===== dataframe =====
+    df = pd.DataFrame(rows.values())
+
+    # ===== split columns =====
+    normal_cols = [
+        c for c in df.columns
+        if not isinstance(c, tuple)
+    ]
+
+    metric_cols = [
+        c for c in df.columns
+        if isinstance(c, tuple)
+    ]
+
+    df = df[normal_cols + metric_cols]
+
+    # ===== build multi columns =====
+    df.columns = pd.MultiIndex.from_tuples([
+        ("", "", c) if not isinstance(c, tuple) else c
+        for c in df.columns
+    ])
+
+    return df
+
+def dict_to_records(data):
+    """
+    Input:
+    {
+        model_name: {
+            (level1, level2, ..., leveln): {
+                metric: value
+            }
+        }
+    }
+
+    Output:
+    [
+        {
+            "model": model_name,
+            "levels": [...],
+            "metrics": {...}
+        },
+        ...
+    ]
+    """
+
+    records = []
+
+    for model_name, model_data in data.items():
+
+        for levels, metrics in model_data.items():
+
+            # ensure tuple/list compatible
+            if not isinstance(levels, (tuple, list)):
+                levels = [levels]
+
+            records.append(
+                {
+                    "model": model_name,
+                    "levels": list(levels),
+                    "metrics": metrics
+                }
+            )
+
+    return records
+
+# %% [code]
+os.makedirs(f'{checkpoints_dir}/results', exist_ok=True)
+testset = KLTNDataset(data_test, range(len(data_test)), label2id, tokenizer, **val_memory_params)
+generator = torch.Generator()
+test_loader = DataLoader(testset, generator=generator, collate_fn=collate_fn, **val_loader_params)
+eval_fn = CustomEvalFn(**eval_func_params)
+analyzer = SpanErrorAnalyzer()
+my_model = IEModel(
+    num_labels=len(label2id), 
+    **model_params
+)
+total_params = sum(p.numel() for p in my_model.parameters())
+print(f"Total params: {total_params:,}")
+
+# %% [code]
+start_time = time.time()
+
+best_score, best_analyze_result, best_pred_test = test(my_model, best_models, test_loader, eval_fn, analyzer, device, id2label, tokenizer)
+last_score, last_analyze_result, last_pred_test = test(my_model, last_models, test_loader, eval_fn, analyzer, device, id2label, tokenizer)
+
+result_test = {"Best model": best_score, "Last model": last_score}
+analyze_result = {"Best model": best_analyze_result, "Last model": last_analyze_result}
+analyze_result_sumary = {"Best model": best_analyze_result['summary'], "Last model": last_analyze_result['summary']}
+pred_test = {"Best model": best_pred_test, "Last model": last_pred_test}
+
+with open(f"{checkpoints_dir}/results/{state_dict_save_name}_test_{my_model.keep_neighbor}.json", "w", encoding="utf-8") as f:
+    json.dump(dict_to_records(result_test), f, ensure_ascii=False, indent=2)
+
+with open(f"{checkpoints_dir}/results/{state_dict_save_name}_error_analyze_result_{my_model.keep_neighbor}.json", "w", encoding="utf-8") as f:
+    json.dump(analyze_result, f, ensure_ascii=False, indent=2)
+
+with open(f"{checkpoints_dir}/results/{state_dict_save_name}_pred_test_{my_model.keep_neighbor}.json", "w", encoding="utf-8") as f:
+    json.dump(pred_test, f, ensure_ascii=False, indent=2)
+    
+print('Test:', time.time() - start_time, 's --> Done!')
+print(json.dumps(analyze_result_sumary, ensure_ascii=False, indent=4))
+
+# %% [code]
+result_test_df = dict_to_df(result_test, row_names=['keep_neightbor', 'alpha'])
+
+col = ("Best model", "Ent-C", "f1")
+sorted_df = result_test_df.sort_values(
+    by=col,
+    ascending=False
+).reset_index(drop=True)
+sorted_df.to_excel(f"{checkpoints_dir}/results/{state_dict_save_name}_test_df_.xlsx")
+
+sorted_df
+
+# %% [code]
+def get_avg_best_score(logs):
+    return float(np.mean([list(log.values())[-1]['best_score'] for log in logs]))
+    
+def get_avg_log(logs, epochs):
+    avg_log = {}
+
+    for epoch in range(1, epochs + 1):
+        val_score = 0.0
+        train_loss = 0.0
+        n_eval = 0
+
+        for idx in range(len(logs)):
+            log = logs[idx].get(epoch, logs[idx].get(str(epoch)))
+            if log is None:
+                continue
+
+            val_score += log.get('val_score', 0.0)
+            train_loss += log.get('train_loss', 0.0)
+            n_eval += 1
+
+        if n_eval == 0:
+            continue
+
+        avg_log[epoch] = {
+            'train_loss': train_loss / n_eval,
+            'val_score': val_score / n_eval if val_score != 0 else float('inf')
+        }
+
+    return avg_log
+
+def parse_label_key(label: str):
+    try:
+        first = float(label.split('_', 1)[0]) # số đầu: trước dấu _
+        last = float(re.findall(r'_(\d+(?:\.\d+)?)$', label)[0])
+        return first, last
+    except:
+        return (0, 0)
+
+def plot_training_logs(logs_dict, save_path=None, figsize=(24, 10)):
+    fig, axes = plt.subplots(1, 2, figsize=figsize)
+
+    # ===== Plot Train Loss =====
+    for name, log in logs_dict.items():
+        epochs = sorted(log.keys())
+        train_loss = [log[e]['train_loss'] for e in epochs]
+        axes[0].plot(epochs, train_loss, label=name)
+
+    axes[0].set_xlabel('Epoch')
+    axes[0].set_ylabel('Train Loss')
+    axes[0].set_title('Training Loss')
+    axes[0].grid(True)
+
+    # ===== Plot Validation Score =====
+    for name, log in logs_dict.items():
+        epochs = sorted(log.keys())
+        val_score = [log[e]['val_score'] for e in epochs]
+        axes[1].plot(epochs, val_score, label=name)
+
+    axes[1].set_xlabel('Epoch')
+    axes[1].set_ylabel('Validation Score')
+    axes[1].set_title('Validation Score')
+    axes[1].grid(True)
+
+    # ===== Shared Legend =====
+    handles, labels = axes[0].get_legend_handles_labels()
+    pairs = list(zip(handles, labels))
+    pairs_sorted = sorted(
+        pairs,
+        key=lambda x: parse_label_key(x[1])
+    )
+    handles_sorted, labels_sorted = zip(*pairs_sorted)
+    
+    axes[0].legend(
+        handles_sorted,
+        labels_sorted,
+        loc='center left',
+        bbox_to_anchor=(1.01, 0.5),
+        borderaxespad=0.
+    )
+
+    plt.tight_layout(rect=[0, 0, 1, 1])
+
+    if save_path is not None:
+        os.makedirs(os.path.dirname(save_path), exist_ok=True) if os.path.dirname(save_path) else None
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+
+    plt.show()
+
+# %% [code]
+if not test_only:
+    snapshot_download(repo_id=repo_name, local_dir="", repo_type="model", allow_patterns=["**/*ent*.json"], ignore_patterns=["**/*crf*.json"])
+    get_ipython().system('rm -rf .cache .gitattributes')
+
+# %% [code]
+if not test_only:
+    experiments = {}
+    for experiment in os.listdir(pretrained_dir):
+        if '.virtual_documents' in experiment:
+            continue
+        experiment_logs = []
+        try:
+            for seed in SEEDS:
+                for fold_idx in range(nfolds):
+                    with open(f"{pretrained_dir}/{experiment}/logs/{experiment}_s{seed}_f{fold_idx}_logging.json", "r", encoding="utf-8") as f:
+                        experiment_log = json.load(f)
+                    experiment_logs.append(experiment_log)
+        except:
+            pass
+        experiments[experiment] = get_avg_log(experiment_logs, 1000)
+    experiments[state_dict_save_name] = get_avg_log(training_logs, 1000)
+
+# %% [code]
+if not test_only:
+    score = get_avg_best_score(training_logs)
+    state_dict_save_name, score
+
+# %% [code]
+if not test_only:
+    plot_training_logs(experiments, save_path=f'{checkpoints_dir}/logs/{state_dict_save_name}_log_plot.jpg', figsize=(18, 7.5))
+

20_entities_top_50_25/lasts/20_entities_top_50_25_s26092004_f0_last_ema.pth

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

20_entities_top_50_25/logs/20_entities_top_50_25_s26092004_f0_logging.json

@@ -0,0 +1 @@
+{"1": {"lr": [2e-05, 0.0005], "train_loss": 0.15836021304130554, "total": 0.15836021632305933, "span_loss": 0.013298947295374947, "start_loss": 0.06792396643496279, "end_loss": 0.0771371789189142}, "2": {"lr": [1.988303923565381e-05, 0.0004969282409784868], "train_loss": 0.1167537197470665, "total": 0.116753720517529, "span_loss": 0.00795133610955159, "start_loss": 0.05067283071706653, "end_loss": 0.05812967469933447}, "3": {"lr": [1.9535036904803962e-05, 0.0004877886008156408], "train_loss": 0.1098012775182724, "total": 0.10980128229097523, "span_loss": 0.007565838455021082, "start_loss": 0.04854550476244803, "end_loss": 0.05368998971579793}, "4": {"lr": [1.8964561979789496e-05, 0.00047280612778499774], "train_loss": 0.10013656318187714, "total": 0.10013656658857079, "span_loss": 0.006966773726674949, "start_loss": 0.04450445884162984, "end_loss": 0.04866513571571278}, "5": {"lr": [1.8185661446562005e-05, 0.00045234974009654937], "train_loss": 0.09494887292385101, "total": 0.09494887683030412, "span_loss": 0.006612358178553999, "start_loss": 0.04236734133315793, "end_loss": 0.04596924288703317}, "6": {"lr": [1.7217514421272206e-05, 0.00042692314190604356], "train_loss": 0.09057747572660446, "total": 0.09057747684689858, "span_loss": 0.006500215887556934, "start_loss": 0.040513803600398704, "end_loss": 0.043563458958173214}, "7": {"lr": [1.60839598967785e-05, 0.00039715242044697206], "train_loss": 0.08300333470106125, "total": 0.08300333465524996, "span_loss": 0.006016988738280158, "start_loss": 0.03738455775065286, "end_loss": 0.03960186279706278}, "8": {"lr": [1.4812909747525698e-05, 0.00036377062968501693], "train_loss": 0.07877770066261292, "total": 0.0787776991799932, "span_loss": 0.0058678082742792325, "start_loss": 0.03547269715617794, "end_loss": 0.037437170294158745, "val_score": 1.6879098478651393, "best_score": 1.6879098478651393, "new_best_model": true, "precision": 0.6927880624134277, "recall": 0.7066723232711557, "f1": 0.698129392068663, "span_recall": 0.9897804557964766}, "9": {"lr": [1.3435661446562005e-05, 0.0003275997400965494], "train_loss": 0.07042290270328522, "total": 0.07042290659307968, "span_loss": 0.005389453446551899, "start_loss": 0.03172466646965287, "end_loss": 0.033308669933065346, "val_score": 1.6880608116033966, "best_score": 1.6880608116033966, "new_best_model": true, "precision": 0.6926097764085044, "recall": 0.7085987154401036, "f1": 0.6989633038308182, "span_recall": 0.9890975077725791}, "10": {"lr": [1.1986127417882198e-05, 0.00028953039902753766], "train_loss": 0.06412826478481293, "total": 0.06412826802492008, "span_loss": 0.005015367366136153, "start_loss": 0.02899997043769542, "end_loss": 0.030112998454913272, "val_score": 1.6844831534157307, "best_score": 1.6880608116033966, "new_best_model": false, "precision": 0.6875223989713674, "recall": 0.7089024167935534, "f1": 0.6965351773006304, "span_recall": 0.9879479761150991}, "11": {"lr": [1.0500000000000003e-05, 0.0002505], "train_loss": 0.05869597569108009, "total": 0.05869597706333628, "span_loss": 0.004821416450786218, "start_loss": 0.026491579894848946, "end_loss": 0.027382992978466497, "val_score": 1.6814694558381817, "best_score": 1.6880608116033966, "new_best_model": false, "precision": 0.6850706713058013, "recall": 0.7074393355117581, "f1": 0.6945111173782986, "span_recall": 0.9869583384598831}, "12": {"lr": [9.013872582117811e-06, 0.00021146960097246246], "train_loss": 0.052141811698675156, "total": 0.05214181114269276, "span_loss": 0.004431201333503872, "start_loss": 0.023577266328490734, "end_loss": 0.024133408249524, "val_score": 1.6788781985189072, "best_score": 1.6880608116033966, "new_best_model": false, "precision": 0.6824626256370582, "recall": 0.7058162357753008, "f1": 0.6924016210918472, "span_recall": 0.9864765774270605}, "13": {"lr": [7.564338553438001e-06, 0.00017340025990345064], "train_loss": 0.04715990647673607, "total": 0.04715990607068151, "span_loss": 0.004110376017929256, "start_loss": 0.021314065505699982, "end_loss": 0.021735484004353864, "val_score": 1.6761510530355455, "best_score": 1.6880608116033966, "new_best_model": false, "precision": 0.6787482509223736, "recall": 0.7043477681110263, "f1": 0.6897907039476386, "span_recall": 0.9863603490879074}, "14": {"lr": [6.1870902524743065e-06, 0.00013722937031498307], "train_loss": 0.043062418699264526, "total": 0.043062417982942636, "span_loss": 0.0039032074644153782, "start_loss": 0.019475716195871472, "end_loss": 0.019683522309185474, "val_score": 1.6738327289989086, "best_score": 1.6880608116033966, "new_best_model": false, "precision": 0.6776761968384113, "recall": 0.7027311676897512, "f1": 0.6884340121848059, "span_recall": 0.985398716814104}, "15": {"lr": [4.916040103221507e-06, 0.00010384757955302797], "train_loss": 0.039497774094343185, "total": 0.03949777636200173, "span_loss": 0.0037452027143347255, "start_loss": 0.017812720352607843, "end_loss": 0.01793987088659211, "val_score": 1.6726570404456254, "best_score": 1.6880608116033966, "new_best_model": false, "precision": 0.6782996836928802, "recall": 0.7014953495122939, "f1": 0.6881792029631572, "span_recall": 0.9844778374824673}, "16": {"lr": [3.7824855787278e-06, 7.40768580939564e-05], "train_loss": 0.036539915949106216, "total": 0.03653991682160294, "span_loss": 0.0035204935367044755, "start_loss": 0.016551774197147172, "end_loss": 0.016467658949671275, "val_score": 1.6698516693773606, "best_score": 1.6880608116033966, "new_best_model": false, "precision": 0.6760791197776373, "recall": 0.6998378552184006, "f1": 0.6861856178380178, "span_recall": 0.9836660515393411}, "17": {"lr": [2.814338553438001e-06, 4.865025990345063e-05], "train_loss": 0.03419886529445648, "total": 0.03419886546937229, "span_loss": 0.0034353670426225317, "start_loss": 0.015404574893718995, "end_loss": 0.015358882752658936, "val_score": 1.669118306472923, "best_score": 1.6880608116033966, "new_best_model": false, "precision": 0.6767968954388489, "recall": 0.699000197096468, "f1": 0.6862313053001838, "span_recall": 0.9828870011727384}, "18": {"lr": [2.0354380202105066e-06, 2.8193872215002235e-05], "train_loss": 0.03241745010018349, "total": 0.03241744875707994, "span_loss": 0.003389500052386519, "start_loss": 0.0145352450477937, "end_loss": 0.014492702324207832, "val_score": 1.6678282196605587, "best_score": 1.6880608116033966, "new_best_model": false, "precision": 0.6749447392504405, "recall": 0.6987369641707758, "f1": 0.6851534395678112, "span_recall": 0.9826747800927468}}

20_entities_top_50_25/r1s/20_entities_top_50_25_s26092004_f0_r1_vs1.68806_ema.pth

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

20_entities_top_50_25/results/20_entities_top_50_25_test_2.json

@@ -0,0 +1,1066 @@
+[
+  {
+    "model": "Best model",
+    "levels": [
+      0,
+      0.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.7274748732538223,
+        "recall": 0.7602007621519098,
+        "f1": 0.7434778606503211
+      },
+      "Ent-C": {
+        "precision": 0.6501823356749531,
+        "recall": 0.6788002600049412,
+        "f1": 0.6641831678165434
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      1,
+      0.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.7274748732538223,
+        "recall": 0.7602007621519098,
+        "f1": 0.7434778606503211
+      },
+      "Ent-C": {
+        "precision": 0.6501823356749531,
+        "recall": 0.6788002600049412,
+        "f1": 0.6641831678165434
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      2,
+      0.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.7274748732538223,
+        "recall": 0.7602007621519098,
+        "f1": 0.7434778606503211
+      },
+      "Ent-C": {
+        "precision": 0.6501823356749531,
+        "recall": 0.6788002600049412,
+        "f1": 0.6641831678165434
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      3,
+      0.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.7274748732538223,
+        "recall": 0.7602007621519098,
+        "f1": 0.7434778606503211
+      },
+      "Ent-C": {
+        "precision": 0.6501823356749531,
+        "recall": 0.6788002600049412,
+        "f1": 0.6641831678165434
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      0,
+      0.2
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.707413647851131,
+        "recall": 0.7804628682955846,
+        "f1": 0.7421450301433067
+      },
+      "Ent-C": {
+        "precision": 0.6323504633524579,
+        "recall": 0.6970006500132816,
+        "f1": 0.6631034889818218
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      1,
+      0.2
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.707413647851131,
+        "recall": 0.7804628682955846,
+        "f1": 0.7421450301433067
+      },
+      "Ent-C": {
+        "precision": 0.6323504633524579,
+        "recall": 0.6970006500132816,
+        "f1": 0.6631034889818218
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      2,
+      0.2
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.707413647851131,
+        "recall": 0.7804628682955846,
+        "f1": 0.7421450301433067
+      },
+      "Ent-C": {
+        "precision": 0.6323504633524579,
+        "recall": 0.6970006500132816,
+        "f1": 0.6631034889818218
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      3,
+      0.2
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.707413647851131,
+        "recall": 0.7804628682955846,
+        "f1": 0.7421450301433067
+      },
+      "Ent-C": {
+        "precision": 0.6323504633524579,
+        "recall": 0.6970006500132816,
+        "f1": 0.6631034889818218
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      0,
+      0.4
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6903215329646879,
+        "recall": 0.7902221396033179,
+        "f1": 0.7369014034728304
+      },
+      "Ent-C": {
+        "precision": 0.6169210782716653,
+        "recall": 0.7055436902212782,
+        "f1": 0.6582629363257605
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      1,
+      0.4
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6903215329646879,
+        "recall": 0.7902221396033179,
+        "f1": 0.7369014034728304
+      },
+      "Ent-C": {
+        "precision": 0.6169210782716653,
+        "recall": 0.7055436902212782,
+        "f1": 0.6582629363257605
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      2,
+      0.4
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6903215329646879,
+        "recall": 0.7902221396033179,
+        "f1": 0.7369014034728304
+      },
+      "Ent-C": {
+        "precision": 0.6169210782716653,
+        "recall": 0.7055436902212782,
+        "f1": 0.6582629363257605
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      3,
+      0.4
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6903215329646879,
+        "recall": 0.7902221396033179,
+        "f1": 0.7369014034728304
+      },
+      "Ent-C": {
+        "precision": 0.6169210782716653,
+        "recall": 0.7055436902212782,
+        "f1": 0.6582629363257605
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      0,
+      0.5
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6829639113381748,
+        "recall": 0.7932893391571769,
+        "f1": 0.7340041229942705
+      },
+      "Ent-C": {
+        "precision": 0.6105465311669916,
+        "recall": 0.7085151824675378,
+        "f1": 0.6558927140161004
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      1,
+      0.5
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6829092654819308,
+        "recall": 0.7932893391571769,
+        "f1": 0.7339725624263866
+      },
+      "Ent-C": {
+        "precision": 0.6104976796282521,
+        "recall": 0.7085151824675378,
+        "f1": 0.6558645241885204
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      2,
+      0.5
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6829092654819308,
+        "recall": 0.7932893391571769,
+        "f1": 0.7339725624263866
+      },
+      "Ent-C": {
+        "precision": 0.6104976796282521,
+        "recall": 0.7085151824675378,
+        "f1": 0.6558645241885204
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      3,
+      0.5
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6829092654819308,
+        "recall": 0.7932893391571769,
+        "f1": 0.7339725624263866
+      },
+      "Ent-C": {
+        "precision": 0.6104976796282521,
+        "recall": 0.7085151824675378,
+        "f1": 0.6558645241885204
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      0,
+      0.6
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6757482429118883,
+        "recall": 0.7953341388597496,
+        "f1": 0.7306805517739423
+      },
+      "Ent-C": {
+        "precision": 0.6041222459128136,
+        "recall": 0.7103723651214501,
+        "f1": 0.6529532213892284
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      1,
+      0.6
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6756948831327576,
+        "recall": 0.7953341388597496,
+        "f1": 0.7306493567709114
+      },
+      "Ent-C": {
+        "precision": 0.604074542008367,
+        "recall": 0.7103723651214501,
+        "f1": 0.6529253567015821
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      2,
+      0.6
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6756948831327576,
+        "recall": 0.7953341388597496,
+        "f1": 0.7306493567709114
+      },
+      "Ent-C": {
+        "precision": 0.604074542008367,
+        "recall": 0.7103723651214501,
+        "f1": 0.6529253567015821
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      3,
+      0.6
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6756948831327576,
+        "recall": 0.7953341388597496,
+        "f1": 0.7306493567709114
+      },
+      "Ent-C": {
+        "precision": 0.604074542008367,
+        "recall": 0.7103723651214501,
+        "f1": 0.6529253567015821
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      0,
+      0.8
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.662070559235355,
+        "recall": 0.7988660656187387,
+        "f1": 0.7240638508201541
+      },
+      "Ent-C": {
+        "precision": 0.5918194423042737,
+        "recall": 0.7134367165004053,
+        "f1": 0.6469622282119929
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      1,
+      0.8
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6619685767092868,
+        "recall": 0.7988660656187387,
+        "f1": 0.7240028590453211
+      },
+      "Ent-C": {
+        "precision": 0.5917282809607273,
+        "recall": 0.7134367165004053,
+        "f1": 0.6469077540628477
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      2,
+      0.8
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6619175972270605,
+        "recall": 0.7988660656187387,
+        "f1": 0.7239723670109903
+      },
+      "Ent-C": {
+        "precision": 0.5916827108197215,
+        "recall": 0.7134367165004053,
+        "f1": 0.6468805204281683
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      3,
+      0.8
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6619175972270605,
+        "recall": 0.7988660656187387,
+        "f1": 0.7239723670109903
+      },
+      "Ent-C": {
+        "precision": 0.5916827108197215,
+        "recall": 0.7134367165004053,
+        "f1": 0.6468805204281683
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      0,
+      1.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6443930117958455,
+        "recall": 0.8022121014956759,
+        "f1": 0.714693821912725
+      },
+      "Ent-C": {
+        "precision": 0.5754815589065436,
+        "recall": 0.7157581948177957,
+        "f1": 0.6380002433720193
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      1,
+      1.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6443233559747373,
+        "recall": 0.8023050469367019,
+        "f1": 0.7146878572878846
+      },
+      "Ent-C": {
+        "precision": 0.5754273344774387,
+        "recall": 0.7158510539504913,
+        "f1": 0.6380038020601749
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      2,
+      1.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6442752649644392,
+        "recall": 0.8023050469367019,
+        "f1": 0.7146582721630242
+      },
+      "Ent-C": {
+        "precision": 0.5753843857287838,
+        "recall": 0.7158510539504913,
+        "f1": 0.6379774022668191
+      }
+    }
+  },
+  {
+    "model": "Best model",
+    "levels": [
+      3,
+      1.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6443018135676958,
+        "recall": 0.802397992377728,
+        "f1": 0.7147114777977568
+      },
+      "Ent-C": {
+        "precision": 0.5754160758261248,
+        "recall": 0.715943913083187,
+        "f1": 0.6380337587131018
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      0,
+      0.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.7241286863264306,
+        "recall": 0.7531369086339315,
+        "f1": 0.7383479835200943
+      },
+      "Ent-C": {
+        "precision": 0.6447145027248282,
+        "recall": 0.6699786423988578,
+        "f1": 0.6571038201378601
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      1,
+      0.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.7241286863264306,
+        "recall": 0.7531369086339315,
+        "f1": 0.7383479835200943
+      },
+      "Ent-C": {
+        "precision": 0.6447145027248282,
+        "recall": 0.6699786423988578,
+        "f1": 0.6571038201378601
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      2,
+      0.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.7240639799832685,
+        "recall": 0.7531369086339315,
+        "f1": 0.7383143457985305
+      },
+      "Ent-C": {
+        "precision": 0.6446568977835555,
+        "recall": 0.6699786423988578,
+        "f1": 0.657073898739702
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      3,
+      0.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.72399928520307,
+        "recall": 0.7531369086339315,
+        "f1": 0.7382807111417676
+      },
+      "Ent-C": {
+        "precision": 0.6445993031353126,
+        "recall": 0.6699786423988578,
+        "f1": 0.6570439800663778
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      0,
+      0.2
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6997393424697723,
+        "recall": 0.7734919602186323,
+        "f1": 0.734769551783911
+      },
+      "Ent-C": {
+        "precision": 0.6235076509159042,
+        "recall": 0.6886433280706763,
+        "f1": 0.6544588045250469
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      1,
+      0.2
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6996805111815203,
+        "recall": 0.7734919602186323,
+        "f1": 0.7347371159230409
+      },
+      "Ent-C": {
+        "precision": 0.6234552332907747,
+        "recall": 0.6886433280706763,
+        "f1": 0.6544299279450505
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      2,
+      0.2
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6996805111815203,
+        "recall": 0.7734919602186323,
+        "f1": 0.7347371159230409
+      },
+      "Ent-C": {
+        "precision": 0.6234552332907747,
+        "recall": 0.6886433280706763,
+        "f1": 0.6544299279450505
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      3,
+      0.2
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.699621689785036,
+        "recall": 0.7734919602186323,
+        "f1": 0.7347046829257585
+      },
+      "Ent-C": {
+        "precision": 0.6234028244782923,
+        "recall": 0.6886433280706763,
+        "f1": 0.6544010539131748
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      0,
+      0.4
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6784908094157702,
+        "recall": 0.7822288316750792,
+        "f1": 0.7266761595974007
+      },
+      "Ent-C": {
+        "precision": 0.604272470777425,
+        "recall": 0.6960720586863255,
+        "f1": 0.646931901471241
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      1,
+      0.4
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.678272082526855,
+        "recall": 0.7822288316750792,
+        "f1": 0.7265506921369767
+      },
+      "Ent-C": {
+        "precision": 0.604077685550323,
+        "recall": 0.6960720586863255,
+        "f1": 0.6468202556164152
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      2,
+      0.4
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.678272082526855,
+        "recall": 0.7822288316750792,
+        "f1": 0.7265506921369767
+      },
+      "Ent-C": {
+        "precision": 0.604077685550323,
+        "recall": 0.6960720586863255,
+        "f1": 0.6468202556164152
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      3,
+      0.4
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.678272082526855,
+        "recall": 0.7822288316750792,
+        "f1": 0.7265506921369767
+      },
+      "Ent-C": {
+        "precision": 0.604077685550323,
+        "recall": 0.6960720586863255,
+        "f1": 0.6468202556164152
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      0,
+      0.5
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.669253920481016,
+        "recall": 0.7853889766699642,
+        "f1": 0.7226854771779272
+      },
+      "Ent-C": {
+        "precision": 0.5957868060500547,
+        "recall": 0.698579255269107,
+        "f1": 0.6431013798830957
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      1,
+      0.5
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6689889953284229,
+        "recall": 0.7853889766699642,
+        "f1": 0.7225309911836113
+      },
+      "Ent-C": {
+        "precision": 0.5955509816334741,
+        "recall": 0.698579255269107,
+        "f1": 0.6429639709296068
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      2,
+      0.5
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6690151994928205,
+        "recall": 0.7854819221109903,
+        "f1": 0.7225856048813833
+      },
+      "Ent-C": {
+        "precision": 0.5955829969123758,
+        "recall": 0.6986721144018027,
+        "f1": 0.6430219589658785
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      3,
+      0.5
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6689622417472738,
+        "recall": 0.7854819221109903,
+        "f1": 0.7225547145937161
+      },
+      "Ent-C": {
+        "precision": 0.5955358556271999,
+        "recall": 0.6986721144018027,
+        "f1": 0.6429944829173719
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      0,
+      0.6
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6573486139068783,
+        "recall": 0.7890138488699795,
+        "f1": 0.7171883530858584
+      },
+      "Ent-C": {
+        "precision": 0.5843592721636978,
+        "recall": 0.7008078744538018,
+        "f1": 0.637307882221732
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      1,
+      0.6
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6570942023371336,
+        "recall": 0.7890138488699795,
+        "f1": 0.7170369069423463
+      },
+      "Ent-C": {
+        "precision": 0.5841331269345323,
+        "recall": 0.7008078744538018,
+        "f1": 0.6371733666477617
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      2,
+      0.6
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6570698862312072,
+        "recall": 0.7891067943110055,
+        "f1": 0.7170608058518843
+      },
+      "Ent-C": {
+        "precision": 0.5841201052464138,
+        "recall": 0.7009007335864974,
+        "f1": 0.6372039965603209
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      3,
+      0.6
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6570190373002189,
+        "recall": 0.7891067943110055,
+        "f1": 0.7170305258463955
+      },
+      "Ent-C": {
+        "precision": 0.5840749052073171,
+        "recall": 0.7009007335864974,
+        "f1": 0.6371771012369758
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      0,
+      0.8
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6334001633618226,
+        "recall": 0.7928246119520468,
+        "f1": 0.7042020919825571
+      },
+      "Ent-C": {
+        "precision": 0.5627737768204301,
+        "recall": 0.7038722258327571,
+        "f1": 0.6254641422683939
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      1,
+      0.8
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6331180880274377,
+        "recall": 0.7928246119520468,
+        "f1": 0.7040277269868673
+      },
+      "Ent-C": {
+        "precision": 0.5625231910942021,
+        "recall": 0.7038722258327571,
+        "f1": 0.6253093499537059
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      2,
+      0.8
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6330513505486546,
+        "recall": 0.7929175573930728,
+        "f1": 0.7040231021388053
+      },
+      "Ent-C": {
+        "precision": 0.5624721768804255,
+        "recall": 0.7039650849654526,
+        "f1": 0.6253144669553269
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      3,
+      0.8
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.6330043778284239,
+        "recall": 0.7929175573930728,
+        "f1": 0.7039940534877868
+      },
+      "Ent-C": {
+        "precision": 0.5624304473621466,
+        "recall": 0.7039650849654526,
+        "f1": 0.625288678664709
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      0,
+      1.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.5998878845206363,
+        "recall": 0.7957059206238537,
+        "f1": 0.684059124143036
+      },
+      "Ent-C": {
+        "precision": 0.5324038394167081,
+        "recall": 0.7056365493539737,
+        "f1": 0.6069004024131198
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      1,
+      1.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.5995517893405703,
+        "recall": 0.7957059206238537,
+        "f1": 0.6838405574435115
+      },
+      "Ent-C": {
+        "precision": 0.5321055948459267,
+        "recall": 0.7056365493539737,
+        "f1": 0.6067065819241425
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      2,
+      1.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.599495868925501,
+        "recall": 0.7957988660648798,
+        "f1": 0.6838384999504639
+      },
+      "Ent-C": {
+        "precision": 0.5320638476613468,
+        "recall": 0.7057294084866693,
+        "f1": 0.6067137619122933
+      }
+    }
+  },
+  {
+    "model": "Last model",
+    "levels": [
+      3,
+      1.0
+    ],
+    "metrics": {
+      "Ent-I": {
+        "precision": 0.5994538962398659,
+        "recall": 0.7957988660648798,
+        "f1": 0.6838111922871924
+      },
+      "Ent-C": {
+        "precision": 0.532026601329694,
+        "recall": 0.7057294084866693,
+        "f1": 0.6066895456687551
+      }
+    }
+  }
+]