config dataclass for models

src/schemas/config.py

@@ -0,0 +1,72 @@
+import json
+from dataclasses import dataclass, fields
+from pathlib import Path
+from transformers import TrainingArguments
+
+VALID_MODES = {"marker", "qa_m", "qa_b"}
+
+
+@dataclass
+class TrainingConfig:
+    mode: str = "marker"
+    data_path: str = "data/data_augmented_256.jsonl"
+    output_dir: str = "models/"
+    model_name: str = "distilbert-base-uncased"
+    max_len: int = 256
+    num_train_epochs: int = 5
+    per_device_train_batch_size: int = 32
+    per_device_eval_batch_size: int = 64
+    gradient_accumulation_steps: int = 1
+    learning_rate: float = 2e-5
+    warmup_ratio: float = 0.1
+    weight_decay: float = 0.01
+    val_split: float = 0.1
+    test_split: float = 0.1
+    early_stopping_patience: int = 3
+    fp16: bool = True
+    seed: int = 42
+    logging_steps: int = 50
+    save_total_limit: int = 2
+    loss_fn: str = "cross_entropy"
+    focal_gamma: float = 2.0
+
+    def __post_init__(self):
+        if self.mode not in VALID_MODES:
+            raise ValueError(f"mode must be one of {VALID_MODES}, got '{self.mode}'")
+        if self.loss_fn not in ("cross_entropy", "focal"):
+            raise ValueError(f"loss_fn must be 'cross_entropy' or 'focal', got '{self.loss_fn}'")
+
+    @classmethod
+    def from_json(cls, path: str | Path) -> "TrainingConfig":
+        with open(path, "r", encoding="utf-8") as f:
+            raw = json.load(f)
+        known = {f.name for f in fields(cls)}
+        unknown = set(raw.keys()) - known
+        if unknown:
+            print(f"Warning: unknown config keys ignored: {unknown}")
+        filtered = {k: v for k, v in raw.items() if k in known}
+        return cls(**filtered)
+
+    def to_training_arguments(self) -> TrainingArguments:
+        return TrainingArguments(
+            output_dir=self.output_dir,
+            num_train_epochs=self.num_train_epochs,
+            per_device_train_batch_size=self.per_device_train_batch_size,
+            per_device_eval_batch_size=self.per_device_eval_batch_size,
+            gradient_accumulation_steps=self.gradient_accumulation_steps,
+            learning_rate=self.learning_rate,
+            warmup_ratio=self.warmup_ratio,
+            weight_decay=self.weight_decay,
+            eval_strategy="epoch",
+            save_strategy="epoch",
+            load_best_model_at_end=True,
+            metric_for_best_model="macro_f1",
+            greater_is_better=True,
+            save_total_limit=self.save_total_limit,
+            logging_dir=f"{self.output_dir}/logs",
+            logging_steps=self.logging_steps,
+            report_to="none",
+            fp16=self.fp16,
+            seed=self.seed,
+            data_seed=self.seed,
+        )