diff --git a/evalharness/model/runner.py b/evalharness/model/runner.py index 9f19207..61f416c 100644 --- a/evalharness/model/runner.py +++ b/evalharness/model/runner.py @@ -38,6 +38,7 @@ async def generate_predictions( attach_context_keys: tuple = ('passage', 'context'), limit_per_task: Optional[int] = None, checkpoint: Union[bool, str] = False, + dataset_name: str = 'adhoc', few_shot_num: int = 0, few_shot_samples: Optional[List[Sample]] = None, few_shot_text: Optional[str] = None, @@ -169,13 +170,11 @@ async def generate_predictions( if checkpoint: from ..eval.checkpoint import CheckpointStore, checkpoint_path - ckpt = checkpoint if isinstance(checkpoint, str) else None - if ckpt is None: - from ..eval.checkpoint import checkpoint_path as _cp - - ckpt = _cp('/data/evalharness' if False else os.path.expanduser('~/.cache/evalharness'), - getattr(spec, 'name', 'adhoc') if spec is not None else 'adhoc', - model_spec) + if isinstance(checkpoint, str): + ckpt = checkpoint + else: + ckpt = checkpoint_path(os.path.expanduser('~/.cache/evalharness'), + dataset_name, adapter.model or str(adapter)) ckpt_store = CheckpointStore(ckpt, model=adapter.model or str(adapter)) restored = ckpt_store.load() else: @@ -254,6 +253,7 @@ async def run_eval( max_input_chars: int = 0, limit_per_task: Optional[int] = None, checkpoint: Union[bool, str] = False, + dataset_name: str = 'adhoc', few_shot_num: int = -1, prompt_style: str = 'strict_letter', ) -> EvalReport: @@ -336,6 +336,7 @@ async def run_eval( system=system, max_turns=max_turns, max_input_chars=max_input_chars, limit_per_task=limit_per_task, checkpoint=checkpoint, + dataset_name=name, few_shot_num=few_shot_num, few_shot_samples=few_shot_samples, few_shot_text=few_shot_text, prompt_style=prompt_style)