wehub-resource-sync
917eedffcf
Main / Python 3.11 - Docs (push) Has been cancelled
Main / Python 3.11 - Build (push) Has been cancelled
Main / Python 3.11 - Lint (push) Has been cancelled
Main / Python 3.11 - Style (push) Has been cancelled
Main / Python 3.11 - Test (push) Has been cancelled
Main / GPU CI (push) Has been cancelled
Main / Release (push) Has been cancelled
Main / Build and Push Docker Images (push) Has been cancelled
520 lines
17 KiB
Python
520 lines
17 KiB
Python
"""Configuration system for OlmOCR training using YAML and dataclasses."""
|
|
|
|
import os
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
import yaml
|
|
from omegaconf import OmegaConf
|
|
|
|
|
|
@dataclass
|
|
class PipelineStepConfig:
|
|
"""Base configuration for pipeline steps."""
|
|
|
|
name: str
|
|
enabled: bool = True
|
|
|
|
|
|
@dataclass
|
|
class FrontMatterParserConfig(PipelineStepConfig):
|
|
"""Configuration for FrontMatterParser step."""
|
|
|
|
name: str = "FrontMatterParser"
|
|
use_page_response_class: bool = True # Whether to use PageResponse dataclass
|
|
|
|
|
|
@dataclass
|
|
class PDFRendererConfig(PipelineStepConfig):
|
|
"""Configuration for PDFRenderer step."""
|
|
|
|
name: str = "PDFRenderer"
|
|
target_longest_image_dim: int = 1024
|
|
|
|
|
|
@dataclass
|
|
class StaticLengthDocumentAnchoringConfig(PipelineStepConfig):
|
|
"""Configuration for StaticLengthDocumentAnchoring step."""
|
|
|
|
name: str = "StaticLengthDocumentAnchoring"
|
|
target_anchor_text_len: int = 6000
|
|
|
|
|
|
@dataclass
|
|
class FinetuningPromptConfig(PipelineStepConfig):
|
|
"""Configuration for FinetuningPrompt step."""
|
|
|
|
name: str = "FinetuningPrompt"
|
|
|
|
|
|
@dataclass
|
|
class NewYamlFinetuningPromptWithAnchoringConfig(PipelineStepConfig):
|
|
"""Configuration for NewYamlFinetuningPromptWithAnchoring step."""
|
|
|
|
name: str = "NewYamlFinetuningPromptWithAnchoring"
|
|
|
|
|
|
@dataclass
|
|
class NewYamlFinetuningPromptWithNoAnchoringConfig(PipelineStepConfig):
|
|
"""Configuration for NewYamlFinetuningPromptWithNoAnchoring step."""
|
|
|
|
name: str = "NewYamlFinetuningPromptWithNoAnchoring"
|
|
|
|
|
|
@dataclass
|
|
class FrontMatterOutputFormatConfig(PipelineStepConfig):
|
|
"""Configuration for FrontMatterOutputFormat step."""
|
|
|
|
name: str = "FrontMatterOutputFormat"
|
|
|
|
|
|
@dataclass
|
|
class JSONOutputFormatConfig(PipelineStepConfig):
|
|
"""Configuration for JSONOutputFormat step."""
|
|
|
|
name: str = "JSONOutputFormat"
|
|
|
|
|
|
@dataclass
|
|
class InstructUserMessagesConfig(PipelineStepConfig):
|
|
"""Configuration for InstructUserMessages step."""
|
|
|
|
name: str = "InstructUserMessages"
|
|
prompt_first: bool = False
|
|
|
|
|
|
@dataclass
|
|
class LatexBracketNormalizerConfig(PipelineStepConfig):
|
|
"""Configuration for LatexBracketNormalizer step."""
|
|
|
|
name: str = "LatexBracketNormalizer"
|
|
|
|
|
|
@dataclass
|
|
class ReformatLatexBoldItalicConfig(PipelineStepConfig):
|
|
"""Configuration for ReformatLatexBoldItalic step."""
|
|
|
|
name: str = "ReformatLatexBoldItalic"
|
|
|
|
|
|
@dataclass
|
|
class TableTransformationConfig(PipelineStepConfig):
|
|
"""Configuration for TableTransformation step."""
|
|
|
|
name: str = "TableTransformation"
|
|
transformation: str = "annotate_dims" # The transformation to apply
|
|
|
|
|
|
@dataclass
|
|
class TokenizerStepConfig(PipelineStepConfig):
|
|
"""Configuration for Tokenizer step."""
|
|
|
|
name: str = "Tokenizer"
|
|
masking_index: int = -100
|
|
end_of_message_token: str = "<|im_end|>"
|
|
|
|
|
|
@dataclass
|
|
class RandomTokenFlipperConfig(PipelineStepConfig):
|
|
"""Configuration for RandomTokenFlipper step."""
|
|
|
|
name: str = "RandomTokenFlipper"
|
|
token_flip_rate: float = 1e-4
|
|
masking_index: int = -100
|
|
|
|
|
|
@dataclass
|
|
class FilterOutRotatedDocumentsConfig(PipelineStepConfig):
|
|
"""Configuration for FilterOutRotatedDocuments step."""
|
|
|
|
name: str = "FilterOutRotatedDocuments"
|
|
|
|
|
|
@dataclass
|
|
class DatasetTextRuleFilterConfig(PipelineStepConfig):
|
|
"""Configuration for DatasetTextRuleFilter step."""
|
|
|
|
name: str = "DatasetTextRuleFilter"
|
|
|
|
|
|
@dataclass
|
|
class RotationAugmentationConfig(PipelineStepConfig):
|
|
"""Configuration for RotationAugmentation step."""
|
|
|
|
name: str = "RotationAugmentation"
|
|
probability: float = 0.5
|
|
|
|
|
|
@dataclass
|
|
class AugraphyBasicAugmentationsConfig(PipelineStepConfig):
|
|
"""Configuration for AugraphyBasicAugmentations step."""
|
|
|
|
name: str = "AugraphyBasicAugmentations"
|
|
probability: float = 0.5 # Overall probability of applying any augmentation
|
|
|
|
|
|
@dataclass
|
|
class DatasetItemConfig:
|
|
"""Configuration for a single dataset item."""
|
|
|
|
root_dir: str
|
|
pipeline: List[Dict[str, Any]] = field(default_factory=list)
|
|
|
|
# Optional sampling
|
|
max_samples: Optional[int] = None
|
|
|
|
|
|
@dataclass
|
|
class DatasetConfig:
|
|
"""Configuration for dataset and data loading."""
|
|
|
|
train: List[Dict[str, Any]] = field(default_factory=list)
|
|
eval: List[Dict[str, Any]] = field(default_factory=list)
|
|
|
|
|
|
@dataclass
|
|
class ModelConfig:
|
|
"""Configuration for model."""
|
|
|
|
name: str = "Qwen/Qwen2.5-VL-7B-Instruct"
|
|
trust_remote_code: bool = False
|
|
|
|
# Model initialization
|
|
load_in_8bit: bool = False
|
|
load_in_4bit: bool = False
|
|
device_map: Any = "auto" # Can be string or dict
|
|
torch_dtype: str = "auto" # "auto", "float16", "bfloat16", "float32"
|
|
|
|
# Flash attention
|
|
use_flash_attention: bool = True
|
|
attn_implementation: Optional[str] = None # "flash_attention_2", "sdpa", "eager"
|
|
|
|
# Model modifications
|
|
freeze_vision_tower: bool = False
|
|
freeze_language_model: bool = False
|
|
|
|
# LoRA configuration (optional)
|
|
use_lora: bool = False
|
|
lora_rank: int = 8
|
|
lora_alpha: int = 32
|
|
lora_dropout: float = 0.1
|
|
lora_target_modules: List[str] = field(default_factory=lambda: ["q_proj", "v_proj", "k_proj", "o_proj"])
|
|
lora_modules_to_save: Optional[List[str]] = None
|
|
|
|
|
|
@dataclass
|
|
class TrainingConfig:
|
|
"""Configuration for training parameters."""
|
|
|
|
output_dir: str = "./outputs"
|
|
num_train_epochs: int = 1
|
|
per_device_train_batch_size: int = 1
|
|
per_device_eval_batch_size: int = 1
|
|
gradient_accumulation_steps: int = 8
|
|
|
|
# Learning rate and scheduler
|
|
learning_rate: float = 2e-5
|
|
lr_scheduler_type: str = "cosine"
|
|
warmup_ratio: float = 0.1
|
|
lr_scheduler_kwargs: Dict[str, Any] = field(default_factory=dict)
|
|
|
|
# Optimization
|
|
optim: str = "adamw_torch" # "adamw_torch", "muon"
|
|
adam_beta1: float = 0.9
|
|
adam_beta2: float = 0.999
|
|
adam_epsilon: float = 1e-8
|
|
weight_decay: float = 0.01
|
|
max_grad_norm: float = 1.0
|
|
|
|
# Muon optimizer specific settings
|
|
muon_momentum: float = 0.95
|
|
muon_lr_multiplier_head: float = 11.0 # Learning rate multiplier for head parameters
|
|
muon_lr_multiplier_embed: float = 30.0 # Learning rate multiplier for embedding parameters
|
|
muon_lr_multiplier_scalar: float = 2.0 # Learning rate multiplier for scalar parameters
|
|
|
|
# Gradient checkpointing
|
|
gradient_checkpointing: bool = False
|
|
gradient_checkpointing_kwargs: Dict[str, Any] = field(default_factory=lambda: {"use_reentrant": False})
|
|
|
|
# Evaluation and checkpointing
|
|
evaluation_strategy: str = "steps"
|
|
eval_steps: int = 500
|
|
save_strategy: str = "steps"
|
|
save_steps: int = 500
|
|
save_total_limit: int = 3
|
|
load_best_model_at_end: bool = True
|
|
metric_for_best_model: str = "eval_loss"
|
|
greater_is_better: bool = False
|
|
|
|
# Logging
|
|
logging_dir: Optional[str] = None
|
|
logging_strategy: str = "steps"
|
|
logging_steps: int = 10
|
|
logging_first_step: bool = True
|
|
report_to: List[str] = field(default_factory=lambda: ["wandb"])
|
|
|
|
# Force seeds to a consistent value for reproducibility
|
|
seed: int = 42
|
|
data_seed: Optional[int] = 42
|
|
|
|
# Performance
|
|
dataloader_drop_last: bool = True
|
|
dataloader_num_workers: int = 16
|
|
|
|
# Data collator settings
|
|
collator_max_token_len: Optional[int] = None
|
|
remove_unused_columns: bool = False # Important for custom datasets
|
|
|
|
# Torch compile settings
|
|
torch_compile: bool = False
|
|
torch_compile_backend: str = "inductor" # "inductor", "aot_eager", "cudagraphs", etc.
|
|
torch_compile_mode: str = "default" # "default", "reduce-overhead", "max-autotune"
|
|
torch_compile_fullgraph: bool = False
|
|
torch_compile_dynamic: bool = False
|
|
|
|
# Early stopping
|
|
use_early_stopping: bool = False
|
|
early_stopping_patience: int = 3
|
|
early_stopping_threshold: float = 0.0
|
|
|
|
|
|
@dataclass
|
|
class Config:
|
|
"""Main configuration class that combines all sub-configs."""
|
|
|
|
model: ModelConfig
|
|
dataset: DatasetConfig
|
|
training: TrainingConfig
|
|
|
|
# Environment
|
|
project_name: str = "olmocr-training"
|
|
run_name: Optional[str] = None
|
|
tags: List[str] = field(default_factory=list)
|
|
notes: Optional[str] = None
|
|
|
|
# Experiment tracking
|
|
experiment_tracker: str = "tensorboard" # "tensorboard", "wandb", "mlflow"
|
|
wandb_project: Optional[str] = None
|
|
wandb_entity: Optional[str] = None
|
|
|
|
# Distributed training
|
|
distributed: bool = False
|
|
local_rank: int = -1
|
|
|
|
@classmethod
|
|
def from_yaml(cls, yaml_path: str | Path) -> "Config":
|
|
"""Load configuration from YAML file."""
|
|
yaml_path = Path(yaml_path)
|
|
if not yaml_path.exists():
|
|
raise FileNotFoundError(f"Config file not found: {yaml_path}")
|
|
|
|
# Load YAML with OmegaConf for better features
|
|
with open(yaml_path, "r") as f:
|
|
yaml_content = yaml.safe_load(f)
|
|
|
|
# Create OmegaConf config for interpolation and validation
|
|
cfg = OmegaConf.create(yaml_content)
|
|
|
|
# Resolve any interpolations
|
|
OmegaConf.resolve(cfg)
|
|
|
|
# Convert to dict and create dataclass
|
|
cfg_dict = OmegaConf.to_container(cfg, resolve=True)
|
|
|
|
# Create sub-configs
|
|
model_cfg = ModelConfig(**cfg_dict.get("model", {}))
|
|
dataset_cfg = DatasetConfig(**cfg_dict.get("dataset", {}))
|
|
training_cfg = TrainingConfig(**cfg_dict.get("training", {}))
|
|
|
|
# Create main config
|
|
main_cfg_dict = {k: v for k, v in cfg_dict.items() if k not in ["model", "dataset", "training"]}
|
|
|
|
return cls(model=model_cfg, dataset=dataset_cfg, training=training_cfg, **main_cfg_dict)
|
|
|
|
def to_yaml(self, yaml_path: str | Path) -> None:
|
|
"""Save configuration to YAML file."""
|
|
yaml_path = Path(yaml_path)
|
|
yaml_path.parent.mkdir(parents=True, exist_ok=True)
|
|
|
|
# Convert to OmegaConf for nice YAML output
|
|
cfg = OmegaConf.structured(self)
|
|
|
|
with open(yaml_path, "w") as f:
|
|
OmegaConf.save(cfg, f)
|
|
|
|
def validate(self) -> None:
|
|
"""Validate configuration values."""
|
|
# Dataset validation - check all train and eval datasets
|
|
for split_name, datasets in [("train", self.dataset.train), ("eval", self.dataset.eval)]:
|
|
for i, dataset_cfg in enumerate(datasets):
|
|
root_dir = dataset_cfg.get("root_dir")
|
|
if not root_dir:
|
|
raise ValueError(f"Missing root_dir for {split_name} dataset {i}")
|
|
if not os.path.exists(root_dir):
|
|
raise ValueError(f"Dataset root directory does not exist: {root_dir}")
|
|
|
|
# Model validation
|
|
if self.model.load_in_8bit and self.model.load_in_4bit:
|
|
raise ValueError("Cannot load in both 8bit and 4bit")
|
|
|
|
# Output directory
|
|
Path(self.training.output_dir).mkdir(parents=True, exist_ok=True)
|
|
|
|
# Logging directory
|
|
if self.training.logging_dir is None:
|
|
self.training.logging_dir = os.path.join(self.training.output_dir, "logs")
|
|
Path(self.training.logging_dir).mkdir(parents=True, exist_ok=True)
|
|
|
|
def to_dict(self) -> Dict[str, Any]:
|
|
"""Convert configuration to dictionary."""
|
|
cfg = OmegaConf.structured(self)
|
|
return OmegaConf.to_container(cfg, resolve=True)
|
|
|
|
def get_pipeline_steps(self, pipeline_config: List[Dict[str, Any]], processor=None):
|
|
"""Create actual pipeline step instances from pipeline configuration.
|
|
|
|
Args:
|
|
pipeline_config: List of pipeline step configurations
|
|
processor: The model processor (required for Tokenizer step)
|
|
|
|
Returns:
|
|
List of initialized pipeline step instances
|
|
"""
|
|
from olmocr.prompts.prompts import PageResponse
|
|
from olmocr.train.dataloader import (
|
|
AugraphyBasicAugmentations,
|
|
DatasetTextRuleFilter,
|
|
FilterOutRotatedDocuments,
|
|
FinetuningPrompt,
|
|
FrontMatterOutputFormat,
|
|
InstructUserMessages,
|
|
JSONOutputFormat,
|
|
LatexBracketNormalizer,
|
|
NewYamlFinetuningPromptWithAnchoring,
|
|
NewYamlFinetuningPromptWithNoAnchoring,
|
|
PDFRenderer,
|
|
RandomTokenFlipper,
|
|
ReformatLatexBoldItalic,
|
|
RotationAugmentation,
|
|
StaticLengthDocumentAnchoring,
|
|
TableTransformation,
|
|
Tokenizer,
|
|
)
|
|
from olmocr.train.front_matter import FrontMatterParser
|
|
|
|
steps = []
|
|
for step_config in pipeline_config:
|
|
if not step_config.get("enabled", True):
|
|
continue
|
|
|
|
step_name = step_config["name"]
|
|
|
|
if step_name == "FrontMatterParser":
|
|
# Handle both old and new config format
|
|
if "front_matter_class" in step_config:
|
|
front_matter_class = PageResponse if step_config["front_matter_class"] == "PageResponse" else None
|
|
else:
|
|
front_matter_class = PageResponse if step_config.get("use_page_response_class", True) else None
|
|
steps.append(FrontMatterParser(front_matter_class=front_matter_class))
|
|
|
|
elif step_name == "PDFRenderer":
|
|
steps.append(PDFRenderer(target_longest_image_dim=step_config.get("target_longest_image_dim", 1024)))
|
|
|
|
elif step_name == "StaticLengthDocumentAnchoring":
|
|
steps.append(StaticLengthDocumentAnchoring(target_anchor_text_len=step_config.get("target_anchor_text_len", 6000)))
|
|
|
|
elif step_name == "FinetuningPrompt":
|
|
steps.append(FinetuningPrompt())
|
|
|
|
elif step_name == "NewYamlFinetuningPromptWithAnchoring":
|
|
steps.append(NewYamlFinetuningPromptWithAnchoring())
|
|
|
|
elif step_name == "NewYamlFinetuningPromptWithNoAnchoring":
|
|
steps.append(NewYamlFinetuningPromptWithNoAnchoring())
|
|
|
|
elif step_name == "JSONOutputFormat":
|
|
steps.append(JSONOutputFormat())
|
|
|
|
elif step_name == "FrontMatterOutputFormat":
|
|
steps.append(FrontMatterOutputFormat())
|
|
|
|
elif step_name == "InstructUserMessages":
|
|
steps.append(InstructUserMessages(prompt_first=step_config.get("prompt_first", False)))
|
|
|
|
elif step_name == "LatexBracketNormalizer":
|
|
steps.append(LatexBracketNormalizer())
|
|
|
|
elif step_name == "Tokenizer":
|
|
if processor is None:
|
|
raise ValueError("Processor must be provided for Tokenizer step")
|
|
steps.append(
|
|
Tokenizer(
|
|
processor=processor,
|
|
masking_index=step_config.get("masking_index", -100),
|
|
end_of_message_token=step_config.get("end_of_message_token", "<|im_end|>"),
|
|
)
|
|
)
|
|
elif step_name == "RandomTokenFlipper":
|
|
if processor is None:
|
|
raise ValueError("Processor must be provided for RandomTokenFlipper step (to get valid tokens)")
|
|
tokenizer = processor.tokenizer
|
|
|
|
# Get all special token IDs to exclude
|
|
special_token_ids = set()
|
|
for token in tokenizer.all_special_tokens:
|
|
special_token_ids.add(tokenizer.convert_tokens_to_ids(token))
|
|
|
|
# Get all token IDs that are not special tokens
|
|
valid_token_ids = []
|
|
for token_id in range(len(tokenizer)):
|
|
if token_id not in special_token_ids:
|
|
valid_token_ids.append(token_id)
|
|
|
|
steps.append(
|
|
RandomTokenFlipper(
|
|
valid_token_ids=valid_token_ids,
|
|
token_flip_rate=step_config.get("token_flip_rate", 1e-4),
|
|
masking_index=step_config.get("masking_index", -100),
|
|
)
|
|
)
|
|
|
|
elif step_name == "FilterOutRotatedDocuments":
|
|
steps.append(FilterOutRotatedDocuments())
|
|
|
|
elif step_name == "DatasetTextRuleFilter":
|
|
steps.append(DatasetTextRuleFilter())
|
|
|
|
elif step_name == "RotationAugmentation":
|
|
steps.append(RotationAugmentation(probability=step_config.get("probability", 0.5)))
|
|
|
|
elif step_name == "AugraphyBasicAugmentations":
|
|
steps.append(AugraphyBasicAugmentations(probability=step_config.get("probability", 0.5)))
|
|
|
|
elif step_name == "ReformatLatexBoldItalic":
|
|
steps.append(ReformatLatexBoldItalic())
|
|
|
|
elif step_name == "TableTransformation":
|
|
steps.append(TableTransformation(transformation=step_config.get("transformation", "annotate_dims")))
|
|
|
|
else:
|
|
raise ValueError(f"Unknown pipeline step: {step_name}")
|
|
|
|
return steps
|
|
|
|
|
|
def create_default_config() -> Config:
|
|
"""Create a default configuration."""
|
|
return Config(model=ModelConfig(), dataset=DatasetConfig(), training=TrainingConfig())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# Example: Create and save default config
|
|
config = create_default_config()
|
|
config.to_yaml("configs/default_config.yaml")
|
|
print("Default config saved to configs/default_config.yaml")
|
|
|
|
# Example: Load from YAML
|
|
# loaded_config = Config.from_yaml("configs/default_config.yaml")
|
|
# print(loaded_config)
|