File size: 1,545 Bytes
d8bfe4a | 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 | from abc import ABC, abstractmethod
class BaseTask(ABC):
def __init__(self, config):
self.config = config
self.dataset = None
self.model = None
self.trainer = None
@abstractmethod
def build_dataset(self):
"""构建并返回训练数据集实例"""
pass
@abstractmethod
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()
|