from .specs import Pipeline
from deep_learning.env.runner import ActionRunner, action


class PipelineRunner(ActionRunner):
    def __init__(self, pipeline: Pipeline):
        self.pipeline = pipeline

    @action
    def train(self):
        return self.pipeline.train()

    @action
    def export(self):
        model_path = self.pipeline.export_model()
        print(f"模型已导出到: {model_path}")
        return model_path

    @action
    def test(self, checkpoint: str = "test"):
        if checkpoint not in ["train", "test"]:
            raise ValueError("test 动作的 checkpoint 只支持 train 或 test")
        return self.pipeline.test_examples(checkpoint=checkpoint)