Skip to content

base

optimus_dl.recipe.eval.base

Evaluation recipe for LLM Baselines models.

EvalRecipe

Recipe for evaluating LLM Baselines models using lm_eval harness.

Orchestrates the evaluation process: 1. Loads a model from a checkpoint or configuration. 2. Wraps it in a LLMBaselinesModel adapter for lm_eval. 3. Runs specified evaluation tasks (e.g., Hellaswag, MMLU). 4. Saves results to JSON.

Parameters:

Name Type Description Default
cfg EvalConfig

Evaluation configuration.

 required
Source code in optimus_dl/recipe/eval/base.py 
 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
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
class EvalRecipe:
    """Recipe for evaluating LLM Baselines models using lm_eval harness.

    Orchestrates the evaluation process:
    1.  Loads a model from a checkpoint or configuration.
    2.  Wraps it in a `LLMBaselinesModel` adapter for `lm_eval`.
    3.  Runs specified evaluation tasks (e.g., Hellaswag, MMLU).
    4.  Saves results to JSON.

    Args:
        cfg: Evaluation configuration.
    """

    def __init__(self, cfg: EvalConfig):
        """Initialize evaluation recipe."""
        self.cfg = cfg
        self.model = None
        self.tokenizer = None

        # Initialize builders using composition
        # ModelBuilder needs a dummy config, but it won't be used since we use build_model_from_checkpoint
        self.model_builder = ModelBuilder(None, [])
        self.checkpoint_manager = CheckpointManager(None)

        # Direct tokenizer build config
        self.tokenizer_config = cfg.common.tokenizer

    def build_eval_model(self, collective: Collective) -> LLMBaselinesModel:
        """Build and load the model for evaluation.

        Loads the model from the configured checkpoint path, initializes the
        tokenizer, and wraps them in the `LLMBaselinesModel` adapter.

        Args:
            collective: Distributed collective for model initialization context.

        Returns:
            An initialized `LLMBaselinesModel` ready for `lm_eval`.
        """
        if self.model is None:
            assert (self.cfg.common.checkpoint_path is not None) ^ (
                self.cfg.common.model is not None
            ), "Either checkpoint_path or model must be specified, but not both"

            device = collective.default_device
            if self.cfg.common.checkpoint_path is not None:
                logger.info(
                    f"Loading model from checkpoint: {self.cfg.common.checkpoint_path}"
                )
                base_model, _ = self.checkpoint_manager.build_model_from_checkpoint(
                    checkpoint_path=self.cfg.common.checkpoint_path, device=device
                )
            else:
                logger.info("Building model from config")
                base_model = self.model_builder.build_model(
                    model_config=self.cfg.common.model,
                    collective=collective,
                )

            # Build tokenizer directly
            self.tokenizer = build_tokenizer(self.tokenizer_config)

            # Wrap in LLMBaselinesModel for lm_eval compatibility
            base_model.eval()
            self.model = LLMBaselinesModel(
                model=base_model.to(device),
                tokenizer=self.tokenizer,
                tokenizer_config=self.tokenizer_config,
                device=device,
            )

        return self.model

    def run_lm_eval(self) -> dict:
        """Run standard benchmarks using the lm_eval harness.

        Sets up the distributed environment (even if single-device), builds the
        model, executes the tasks, and saves results.

        Returns:
            Dictionary containing evaluation metrics and results.
        """
        try:
            from lm_eval import evaluator
        except ImportError as err:
            raise ImportError(
                "lm_eval is required for evaluation. Install with: pip install lm_eval"
            ) from err

        # Build model
        collective = build_best_collective(
            device=None if self.cfg.common.use_gpu else torch.device("cpu"),
            config=DistributedConfig(),
        )
        model = self.build_eval_model(collective=collective)

        logger.info(f"Running lm_eval on tasks: {self.cfg.lm_eval.tasks}")
        logger.info(f"Few-shot examples: {self.cfg.lm_eval.num_fewshot}")

        # Convert tasks to proper format for lm_eval
        raw_tasks = self.cfg.lm_eval.tasks
        if isinstance(raw_tasks, str):
            tasks = [raw_tasks]
        else:
            # Convert to list of strings, handling any nested structure
            tasks = []
            for task in raw_tasks:
                if isinstance(task, str):
                    tasks.append(task)
                else:
                    tasks.append(str(task))

        # Run evaluation
        results = evaluator.simple_evaluate(
            model=model,
            tasks=tasks,
            num_fewshot=self.cfg.lm_eval.num_fewshot,
            batch_size=self.cfg.lm_eval.batch_size,
            limit=self.cfg.lm_eval.limit,
            device=collective.default_device,
            use_cache=None,  # Disable caching for now
            # verbosity="INFO",
        )

        if results is None:
            raise RuntimeError("Evaluation returned no results")

        # Save results if output path is specified
        if self.cfg.lm_eval.output_path:
            import json

            output_path = Path(self.cfg.lm_eval.output_path)
            output_path.parent.mkdir(parents=True, exist_ok=True)

            with open(output_path, "w") as f:
                json.dump(results, f, indent=2)
            logger.info(f"Results saved to: {output_path}")

        return results

__init__(cfg)

Initialize evaluation recipe.

Source code in optimus_dl/recipe/eval/base.py 
33
34
35
36
37
38
39
40
41
42
43
44
45
def __init__(self, cfg: EvalConfig):
    """Initialize evaluation recipe."""
    self.cfg = cfg
    self.model = None
    self.tokenizer = None

    # Initialize builders using composition
    # ModelBuilder needs a dummy config, but it won't be used since we use build_model_from_checkpoint
    self.model_builder = ModelBuilder(None, [])
    self.checkpoint_manager = CheckpointManager(None)

    # Direct tokenizer build config
    self.tokenizer_config = cfg.common.tokenizer

build_eval_model(collective)

Build and load the model for evaluation.

Loads the model from the configured checkpoint path, initializes the tokenizer, and wraps them in the LLMBaselinesModel adapter.

Parameters:

Name Type Description Default
collective Collective

Distributed collective for model initialization context.

 required

Returns:

Type Description
LLMBaselinesModel

An initialized LLMBaselinesModel ready for lm_eval.
Source code in optimus_dl/recipe/eval/base.py 
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def build_eval_model(self, collective: Collective) -> LLMBaselinesModel:
    """Build and load the model for evaluation.

    Loads the model from the configured checkpoint path, initializes the
    tokenizer, and wraps them in the `LLMBaselinesModel` adapter.

    Args:
        collective: Distributed collective for model initialization context.

    Returns:
        An initialized `LLMBaselinesModel` ready for `lm_eval`.
    """
    if self.model is None:
        assert (self.cfg.common.checkpoint_path is not None) ^ (
            self.cfg.common.model is not None
        ), "Either checkpoint_path or model must be specified, but not both"

        device = collective.default_device
        if self.cfg.common.checkpoint_path is not None:
            logger.info(
                f"Loading model from checkpoint: {self.cfg.common.checkpoint_path}"
            )
            base_model, _ = self.checkpoint_manager.build_model_from_checkpoint(
                checkpoint_path=self.cfg.common.checkpoint_path, device=device
            )
        else:
            logger.info("Building model from config")
            base_model = self.model_builder.build_model(
                model_config=self.cfg.common.model,
                collective=collective,
            )

        # Build tokenizer directly
        self.tokenizer = build_tokenizer(self.tokenizer_config)

        # Wrap in LLMBaselinesModel for lm_eval compatibility
        base_model.eval()
        self.model = LLMBaselinesModel(
            model=base_model.to(device),
            tokenizer=self.tokenizer,
            tokenizer_config=self.tokenizer_config,
            device=device,
        )

    return self.model

run_lm_eval()

Run standard benchmarks using the lm_eval harness.

Sets up the distributed environment (even if single-device), builds the model, executes the tasks, and saves results.

Returns:

Type Description
dict

Dictionary containing evaluation metrics and results.
Source code in optimus_dl/recipe/eval/base.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
def run_lm_eval(self) -> dict:
    """Run standard benchmarks using the lm_eval harness.

    Sets up the distributed environment (even if single-device), builds the
    model, executes the tasks, and saves results.

    Returns:
        Dictionary containing evaluation metrics and results.
    """
    try:
        from lm_eval import evaluator
    except ImportError as err:
        raise ImportError(
            "lm_eval is required for evaluation. Install with: pip install lm_eval"
        ) from err

    # Build model
    collective = build_best_collective(
        device=None if self.cfg.common.use_gpu else torch.device("cpu"),
        config=DistributedConfig(),
    )
    model = self.build_eval_model(collective=collective)

    logger.info(f"Running lm_eval on tasks: {self.cfg.lm_eval.tasks}")
    logger.info(f"Few-shot examples: {self.cfg.lm_eval.num_fewshot}")

    # Convert tasks to proper format for lm_eval
    raw_tasks = self.cfg.lm_eval.tasks
    if isinstance(raw_tasks, str):
        tasks = [raw_tasks]
    else:
        # Convert to list of strings, handling any nested structure
        tasks = []
        for task in raw_tasks:
            if isinstance(task, str):
                tasks.append(task)
            else:
                tasks.append(str(task))

    # Run evaluation
    results = evaluator.simple_evaluate(
        model=model,
        tasks=tasks,
        num_fewshot=self.cfg.lm_eval.num_fewshot,
        batch_size=self.cfg.lm_eval.batch_size,
        limit=self.cfg.lm_eval.limit,
        device=collective.default_device,
        use_cache=None,  # Disable caching for now
        # verbosity="INFO",
    )

    if results is None:
        raise RuntimeError("Evaluation returned no results")

    # Save results if output path is specified
    if self.cfg.lm_eval.output_path:
        import json

        output_path = Path(self.cfg.lm_eval.output_path)
        output_path.parent.mkdir(parents=True, exist_ok=True)

        with open(output_path, "w") as f:
            json.dump(results, f, indent=2)
        logger.info(f"Results saved to: {output_path}")

    return results