| from abc import ABC, abstractmethod | |
| class BaseTask(ABC): | |
| def __init__(self, config): | |
| self.config = config | |
| self.dataset = None | |
| self.model = None | |
| self.trainer = None | |
| def build_dataset(self): | |
| """构建并返回训练数据集实例""" | |
| pass | |
| def build_model(self): | |
| """构建并返回模型实例""" | |
| pass | |
| def build_collator(self): | |
| """可选：构建数据整理器（collator）""" | |
| return None | |
| def build_sampler(self): | |
| """可选：构建采样器（支持动态批等）""" | |
| return None | |
| def build_trainer(self): | |
| """构建训练器（默认使用 HuggingFace Trainer）""" | |
| from transformers import Trainer | |
| trainer_args = self.config.get("trainer", {}).get("args", {}) | |
| return Trainer( | |
| model=self.model, | |
| args=trainer_args, | |
| train_dataset=self.dataset, | |
| data_collator=self.build_collator(), | |
| # **其余参数如 eval_dataset 可从 config 中补充 | |
| ) | |
| def run(self): | |
| """执行完整的训练流程：构建各组件并开始训练""" | |
| # 按序构建数据集、模型、collator、sampler 和训练器 | |
| self.dataset = self.build_dataset() | |
| self.model = self.build_model() | |
| self.collator = self.build_collator() | |
| self.sampler = self.build_sampler() | |
| self.trainer = self.build_trainer() | |
| # 开始训练 | |
| self.trainer.train() | |