""" 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