File size: 8,909 Bytes
be90b31 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 | """
Training pipeline and orchestration for multimodal table extraction.
"""
import os
from typing import Any, Dict, Optional
import torch
from torch.utils.data import Dataset
from transformers import TrainingArguments, Trainer
from ..utils.logging import setup_logger
from ..data.collator import QwenVLDataCollator
from ..data.preprocessor import format_qwen_vl_conversation
logger = setup_logger(__name__)
class LazyMultimodalDataset(Dataset):
"""Dataset wrapper that formats multimodal messages on demand."""
def __init__(self, dataset: Any, max_table_chars: Optional[int] = None):
self.dataset = dataset
self.max_table_chars = max_table_chars
def __len__(self) -> int:
return len(self.dataset)
def __getitem__(self, idx: int) -> Dict[str, Any]:
item = self.dataset[idx]
table = item["table"]
if self.max_table_chars and len(table) > self.max_table_chars:
table = table[:self.max_table_chars]
msgs = format_qwen_vl_conversation(
image=item["image"],
target_table=table,
system_prompt=item["system_prompt"]
)
return {"messages": msgs}
def build_training_arguments(cfg: Dict[str, Any]) -> TrainingArguments:
"""Builds HuggingFace TrainingArguments from configuration dictionary."""
t_cfg = cfg.get("training", {})
output_dir = t_cfg.get("output_dir", "outputs/checkpoints")
os.makedirs(output_dir, exist_ok=True)
bf16 = t_cfg.get("bf16", False)
fp16 = t_cfg.get("fp16", True)
# Check CUDA device capability: sm_80+ supports native bf16, older GPUs (e.g. GTX 1660 / Turing) should use fp16
if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] < 8:
bf16 = False
fp16 = True
warmup_steps = t_cfg.get("warmup_steps", 20)
optim = t_cfg.get("optim", "paged_adamw_8bit")
return TrainingArguments(
output_dir=output_dir,
num_train_epochs=t_cfg.get("num_train_epochs", 3),
max_steps=t_cfg.get("max_steps", -1),
per_device_train_batch_size=t_cfg.get("per_device_train_batch_size", 1),
per_device_eval_batch_size=t_cfg.get("per_device_eval_batch_size", 1),
gradient_accumulation_steps=t_cfg.get("gradient_accumulation_steps", 8),
learning_rate=float(t_cfg.get("learning_rate", 2e-4)),
lr_scheduler_type=t_cfg.get("lr_scheduler_type", "cosine"),
warmup_steps=warmup_steps,
weight_decay=t_cfg.get("weight_decay", 0.01),
logging_steps=t_cfg.get("logging_steps", 10),
eval_strategy=t_cfg.get("eval_strategy", "epoch"),
save_strategy=t_cfg.get("save_strategy", "epoch"),
save_total_limit=t_cfg.get("save_total_limit", 2),
load_best_model_at_end=t_cfg.get("load_best_model_at_end", False),
bf16=bf16,
fp16=fp16,
optim=optim,
dataloader_num_workers=t_cfg.get("dataloader_num_workers", 2),
dataloader_pin_memory=False,
gradient_checkpointing=t_cfg.get("gradient_checkpointing", True),
report_to=t_cfg.get("report_to", "none"),
remove_unused_columns=False,
)
from .icdar_loss import ICDARMetricLoss
class ICDARTrainer(Trainer):
"""Custom Hugging Face Trainer that integrates ICDAR Metric Loss (SCST / Token-Weighted)."""
def __init__(self, *args, icdar_loss_module: Optional[ICDARMetricLoss] = None, **kwargs):
super().__init__(*args, **kwargs)
self.icdar_loss_module = icdar_loss_module
def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
if self.icdar_loss_module is None:
# Handle newer transformers versions supporting num_items_in_batch
kwargs = {}
if num_items_in_batch is not None:
kwargs["num_items_in_batch"] = num_items_in_batch
return super().compute_loss(model, inputs, return_outputs=return_outputs, **kwargs)
outputs = model(**inputs)
loss, loss_stats = self.icdar_loss_module(model, inputs, sft_loss=outputs.loss)
if self.state.global_step % max(1, self.args.logging_steps) == 0:
for k, v in loss_stats.items():
if isinstance(v, (int, float)):
self.log({f"train/{k}": v})
return (loss, outputs) if return_outputs else loss
class SciImageTableTrainer:
"""Orchestrator for fine-tuning VLM on scientific table extraction."""
def __init__(
self,
model: Any,
processor: Any,
train_dataset: Any,
eval_dataset: Optional[Any],
cfg: Dict[str, Any],
):
self.model = model
self.processor = processor
self.train_dataset = train_dataset
self.eval_dataset = eval_dataset
self.cfg = cfg
self.training_args = build_training_arguments(cfg)
if self.eval_dataset is None:
self.training_args.eval_strategy = "no"
self.data_collator = QwenVLDataCollator(processor=self.processor)
# Build ICDAR metric loss module if configured
t_cfg = cfg.get("training", {})
loss_type = t_cfg.get("loss_type", "ce")
if loss_type in ("scst", "token_weighted", "hybrid", "icdar_metric"):
m_cfg = t_cfg.get("metric_loss", {})
actual_type = "scst" if loss_type == "icdar_metric" else loss_type
vocab_sz = getattr(getattr(self.model, "config", None), "vocab_size", None)
reward_mode = m_cfg.get("reward_mode", "gpu_tensor")
focal_gamma = float(m_cfg.get("focal_gamma", 0.0))
self.icdar_loss_module = ICDARMetricLoss(
processor=self.processor,
loss_type=actual_type,
lambda_metric=m_cfg.get("lambda_metric", 0.3),
rel_tol=m_cfg.get("rel_tol", 0.05),
temperature=m_cfg.get("temperature", 0.7),
max_gen_tokens=m_cfg.get("max_gen_tokens", 512),
numeric_weight=m_cfg.get("numeric_weight", 3.0),
structural_weight=m_cfg.get("structural_weight", 2.0),
reward_mode=reward_mode,
focal_gamma=focal_gamma,
vocab_size=vocab_sz,
)
logger.info(f"Initialized ICDARMetricLoss (type={actual_type}, lambda={m_cfg.get('lambda_metric', 0.3)}, reward_mode={reward_mode})")
if hasattr(self.icdar_loss_module, "weighted_module") and self.icdar_loss_module.weighted_module is not None:
weights = getattr(self.icdar_loss_module.weighted_module, "token_weights", None)
if weights is not None:
model_dev = next(self.model.parameters()).device
weights = weights.to(model_dev)
setattr(self.model, "token_weights", weights)
setattr(self.model, "focal_gamma", focal_gamma)
if hasattr(self.model, "base_model"):
setattr(self.model.base_model, "token_weights", weights)
setattr(self.model.base_model, "focal_gamma", focal_gamma)
if hasattr(self.model.base_model, "model"):
setattr(self.model.base_model.model, "token_weights", weights)
setattr(self.model.base_model.model, "focal_gamma", focal_gamma)
else:
self.icdar_loss_module = None
def train(self):
"""Executes model training."""
logger.info("Initializing lazy multimodal datasets...")
max_table_chars = self.cfg.get("data", {}).get("max_table_chars", 1500)
train_formatted = LazyMultimodalDataset(self.train_dataset, max_table_chars=max_table_chars)
eval_formatted = (
LazyMultimodalDataset(self.eval_dataset, max_table_chars=max_table_chars)
if self.eval_dataset is not None
else None
)
trainer_cls = ICDARTrainer if self.icdar_loss_module is not None else Trainer
trainer_kwargs = {
"model": self.model,
"args": self.training_args,
"train_dataset": train_formatted,
"eval_dataset": eval_formatted,
"data_collator": self.data_collator,
}
if self.icdar_loss_module is not None:
trainer_kwargs["icdar_loss_module"] = self.icdar_loss_module
trainer = trainer_cls(**trainer_kwargs)
logger.info("Starting training loop...")
train_result = trainer.train()
logger.info("Saving best model adapter and tokenizer/processor...")
output_dir = self.training_args.output_dir
final_adapter_dir = os.path.join(output_dir, "final_adapter")
os.makedirs(final_adapter_dir, exist_ok=True)
trainer.save_model(final_adapter_dir)
self.processor.save_pretrained(final_adapter_dir)
return train_result
|