API Reference

Public API

steb.core.get_model(model_name_or_path, truncate=False, max_tokens=None)

Loads a STEB model.

Dispatches in four stages

1. Match the prefix of model_name_or_path (the part before ":") against each registered class's supported_models list.
2. If the path points at a LISA checkpoint directory (detected via :func:is_lisa_model), route to :class:LISAModel.
3. If nothing matched, inspect the HuggingFace config and route auto-regressive LMs to :class:CausalModel.
4. Otherwise fall back to :class:HFModel.

Parameters:

Name	Type	Description	Default
model_name_or_path	str	The name or path of the model to load.	required
truncate	bool	If True, truncate each text to the token cap instead of chunking and mean-pooling. No-op for non-tokenizer models.	False
max_tokens	Optional[int]	Optional per-text token cap, capped at the model's native maximum. No-op for non-tokenizer models.	None

Returns:

Type	Description
	An instance of a STEBModel.

Source code in steb/core.py

def get_model(
    model_name_or_path: str,
    truncate: bool = False,
    max_tokens: Optional[int] = None,
):
    """
    Loads a STEB model.

    Dispatches in four stages:
      1. Match the prefix of ``model_name_or_path`` (the part before ``":"``)
         against each registered class's ``supported_models`` list.
      2. If the path points at a LISA checkpoint directory (detected via
         :func:`is_lisa_model`), route to :class:`LISAModel`.
      3. If nothing matched, inspect the HuggingFace config and route
         auto-regressive LMs to :class:`CausalModel`.
      4. Otherwise fall back to :class:`HFModel`.

    Args:
        model_name_or_path: The name or path of the model to load.
        truncate: If True, truncate each text to the token cap instead of
            chunking and mean-pooling. No-op for non-tokenizer models.
        max_tokens: Optional per-text token cap, capped at the model's
            native maximum. No-op for non-tokenizer models.

    Returns:
        An instance of a STEBModel.
    """
    kwargs = {"truncate": truncate, "max_tokens": max_tokens}

    registry = get_model_registry()
    # Allow models to be referenced with prefixes, e.g. "lftk:config.yaml" or
    # "tfidfngrams:/path/to/vectorizers.pkl" by matching on the part before ":".
    prefix = model_name_or_path.split(":", 1)[0]
    for model_cls in registry.values():
        if prefix in getattr(model_cls, "supported_models", []):
            return model_cls(model_name_or_path, **kwargs)

    if is_lisa_model(model_name_or_path):
        return registry["lisa"](model_name_or_path, **kwargs)

    if _is_causal_model(model_name_or_path):
        return registry["causal"](model_name_or_path, **kwargs)

    return registry["hf"](model_name_or_path, **kwargs)

steb.core.get_all_datasets()

Retrieves a list of all available STEB datasets.

Returns:

Type	Description
List[str]	A list of all available dataset names.

Source code in steb/core.py

def get_all_datasets() -> List[str]:
    """
    Retrieves a list of all available STEB datasets.

    Returns:
        A list of all available dataset names.
    """
    return DATASET_REGISTRY

steb.core.get_supported_datasets(task_name)

Retrieves a list of datasets that support the given task.

Parameters:

Name	Type	Description	Default
task_name	str	The name of the task.	required

Returns:

Type	Description
List[str]	A list of supported dataset names.

Source code in steb/core.py

def get_supported_datasets(task_name: str) -> List[str]:
    """
    Retrieves a list of datasets that support the given task.

    Args:
        task_name: The name of the task.

    Returns:
        A list of supported dataset names.
    """
    supported_datasets = []
    package_dir = os.path.dirname(os.path.abspath(__file__))
    for dataset_name in DATASET_REGISTRY:
        config_path = os.path.join(package_dir, "steb_datasets", dataset_name, "config.json")
        with open(config_path) as f:
            config = json.load(f)
        if task_name in config.get("tasks", {}):
            supported_datasets.append(dataset_name)
    return supported_datasets

steb.core.get_supported_tasks()

Returns the list of all supported task names.

Source code in steb/core.py

def get_supported_tasks() -> list[str]:
    """Returns the list of all supported task names."""
    return SUPPORTED_TASKS

steb.core.evaluate(model, datasets, episode_sizes=None, task_name=None, n_episodes_per_class=None, batch_size=32, force_reload=False, force_rerun=False, force_rerun_oa=False, progress_bar=False, output_folder=RESULTS_DIR, seed=42, run_name=None)

Evaluates a model on a list of datasets for a given task.

Individual dataset/task failures are caught and logged rather than crashing the entire run. A summary of successes and failures is printed at the end, and a JSON log is written to {output_folder}/logs/.

Parameters:

Name	Type	Description	Default
model		The model to evaluate.	required
datasets	List[str]	A list of dataset names to evaluate on.	required
episode_sizes	Optional[List[int]]	A list of episode sizes to evaluate. If None, uses per-task defaults from TASK_DEFAULTS.	None
task_name	Optional[str]	The name of the task to evaluate. If None, runs all tasks.	None
n_episodes_per_class	Optional[int]	The number of episodes per class. If None, uses per-task defaults from TASK_DEFAULTS.	None
batch_size	int	The batch size for embedding.	32
force_reload	bool	Whether to force reload the datasets.	False
force_rerun	bool	Whether to re-run evaluations even if metrics already exist.	False
force_rerun_oa	bool	Whether to re-run the order_alignment task only, even if its metrics file already exists. Ignored when force_rerun is also set.	False
progress_bar	bool	Whether to show a progress bar.	False
output_folder	str	The folder to save the results to.	RESULTS_DIR
seed	int	The random seed to use.	42
run_name	Optional[str]	An optional label for this run (e.g. preset name). Used in the log filename alongside the model name.	None

Returns:

Type	Description
Dict[str, Any]	A dictionary with "successes", "failures", and "log_path" keys.

Source code in steb/core.py

def evaluate(
    model,
    datasets: List[str],
    episode_sizes: Optional[List[int]] = None,
    task_name: Optional[str] = None,
    n_episodes_per_class: Optional[int] = None,
    batch_size: int = 32,
    force_reload: bool = False,
    force_rerun: bool = False,
    force_rerun_oa: bool = False,
    progress_bar: bool = False,
    output_folder: str = RESULTS_DIR,
    seed: int = 42,
    run_name: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Evaluates a model on a list of datasets for a given task.

    Individual dataset/task failures are caught and logged rather than
    crashing the entire run. A summary of successes and failures is
    printed at the end, and a JSON log is written to
    ``{output_folder}/logs/``.

    Args:
        model: The model to evaluate.
        datasets: A list of dataset names to evaluate on.
        episode_sizes: A list of episode sizes to evaluate. If None,
            uses per-task defaults from TASK_DEFAULTS.
        task_name: The name of the task to evaluate. If None, runs all tasks.
        n_episodes_per_class: The number of episodes per class. If None,
            uses per-task defaults from TASK_DEFAULTS.
        batch_size: The batch size for embedding.
        force_reload: Whether to force reload the datasets.
        force_rerun: Whether to re-run evaluations even if metrics already exist.
        force_rerun_oa: Whether to re-run the order_alignment task only,
            even if its metrics file already exists. Ignored when
            ``force_rerun`` is also set.
        progress_bar: Whether to show a progress bar.
        output_folder: The folder to save the results to.
        seed: The random seed to use.
        run_name: An optional label for this run (e.g. preset name).
            Used in the log filename alongside the model name.

    Returns:
        A dictionary with "successes", "failures", and "log_path" keys.
    """
    set_seed(seed)

    successes: List[Tuple[str, int, str]] = []
    failures: List[Tuple[str, int, str, str]] = []

    def safe_load(
        loader: DatasetLoader,
        dataset_name: str,
        episode_size: int,
        current_task_name: str,
    ) -> Optional[Dict]:
        """
        Attempts to load a dataset, logging and recording the failure if it occurs.

        Args:
            loader: The DatasetLoader to call load() on.
            dataset_name: Name of the dataset (for error reporting).
            episode_size: Current episode size (for error reporting).
            current_task_name: Current task name (for error reporting).

        Returns:
            The loaded dataset, or None if loading failed.
        """
        try:
            return loader.load()
        except Exception as e:
            error_msg = f"{type(e).__name__}: {e}"
            print(colored(f"    FAILED to load dataset: {error_msg}", "red"))
            traceback.print_exc()
            failures.append((dataset_name, episode_size, current_task_name, error_msg))
            return None

    def extract_features(
        dataset,
        episode_size,
        n_episodes_per_class,
        batch_size,
        show_progress=False,
    ):
        """
        Extracts features from the dataset using the specified model.

        Expects dataset format:
            Order Alignment: {"label": [[seq1_most, ..., seq1_least], [seq2_most, ..., seq2_least], ...]}
                Each label maps to a list of ordered sequences. Sequences are grouped into
                episodes, then organized by position (most X, ..., least X).
            Others: {"label": [[text_1, ..., text_N], [text_1, ..., text_M], ...]}

        """
        episodes_by_label = {}
        for label, text_list in dataset.items():
            # Validate nested list format
            assert text_list and isinstance(text_list[0], list), \
                f"Dataset for label '{label}' must be a list of lists"

            seq_len = len(text_list[0])
            if episode_size == -1:
                if current_task_name == "pre_defined_pair_classification":
                    # Keep each text list as its own episode so pairs remain separate
                    episodes_by_label[label] = [[lst] for lst in text_list]
                else:
                    # Group all sequences into a single large episode
                    episodes_by_label[label] = [[[sublist for lst in text_list for sublist in lst]]]
            else:
                # Group sequences into episodes, organize by position
                episodes_by_label[label] = [
                    [[seq[pos] for seq in text_list[i:i+episode_size]] for pos in range(seq_len)]
                    for i in range(0, len(text_list), episode_size)
                ]
                assert len(episodes_by_label[label]) == n_episodes_per_class
                assert all(len(episode[0]) == episode_size for episode in episodes_by_label[label])
        all_episodes = [episode for label, episodes in episodes_by_label.items() for episode in episodes]
        y = [label for label, episodes in episodes_by_label.items() for _ in episodes]
        num_positions = len(all_episodes[0])
        if task_name == "order_alignment":
            assert all(len(episode) == num_positions for episode in all_episodes), \
                ("All entries must have the same number of positions, "
                "functionality for variable-length text sets not implemented.")

        # Flatten episodes for embedding: [[[pos0s], [pos1s], ...], ...] -> [[pos0s], [pos1s], [pos0s], [pos1s], ...]
        flat_episodes = [position for episode in all_episodes for position in episode]
        X_flat = model.embed_multiple(flat_episodes, batch_size, show_progress=show_progress)
        X = [X_flat[i:i+num_positions] for i in range(0, len(X_flat), num_positions)]
        return X, y

    # Cache default embeddings by (dataset, episode_size, n_episodes_per_class)
    # so tasks sharing the same parameters reuse the same embeddings.
    default_cache: Dict[Tuple[str, int, int], Tuple[Any, Any]] = {}

    for dataset_name, current_task_name, task_config, episode_size, resolved_n_episodes in _iter_task_configs(
        datasets, task_name, episode_sizes, n_episodes_per_class,
    ):
        if current_task_name is None:
            error_msg = task_config  # sentinel: error string stored in task_config slot
            print(colored(f"--- Skipping {dataset_name}: {error_msg} ---", "red"))
            failures.append((dataset_name, -1, "config", error_msg))
            continue

        print(colored(f"--- Evaluating {dataset_name} | {current_task_name} (episode size: {episode_size}) ---", "cyan"))

        model_str = os.path.basename(model.model_name_or_path)
        if model_str == "":
            model_str = os.path.basename(os.path.dirname(model.model_name_or_path))
        dset_str = os.path.basename(dataset_name)

        tokens_suffix = (
            f"tokens_{model.effective_max_tokens}"
            if getattr(model, "effective_max_tokens", None) is not None
            else None
        )

        # When n_episodes is not "auto", we can check for existing results early
        if resolved_n_episodes != "auto":
            scores_path = os.path.join(
                output_folder, dset_str, model_str,
                f"{episode_size}_{resolved_n_episodes}", current_task_name,
            )
            if tokens_suffix is not None:
                scores_path = os.path.join(scores_path, tokens_suffix)
            metrics_path = os.path.join(scores_path, "metrics.json")
            if (
                not force_rerun
                and not (force_rerun_oa and current_task_name == "order_alignment")
                and os.path.exists(metrics_path)
            ):
                print(colored(f"    -> Skipping (results already exist)", "yellow"))
                successes.append((dataset_name, episode_size, current_task_name))
                continue

        try:
            if "record_handler" in task_config:
                task_loader = DatasetLoader(
                    dataset_name=dataset_name,
                    episode_size=episode_size,
                    n_episodes_per_class=resolved_n_episodes,
                    force_reload=force_reload,
                    seed=seed,
                    task_name=current_task_name,
                )
                task_dataset = safe_load(task_loader, dataset_name, episode_size, current_task_name)
                if task_dataset is None:
                    continue
                actual_n_episodes = task_loader.n_episodes_per_class
                current_X, current_y = extract_features(
                    task_dataset, episode_size, actual_n_episodes, batch_size, show_progress=progress_bar,
                )
            else:
                dset_loader = DatasetLoader(
                    dataset_name=dataset_name,
                    episode_size=episode_size,
                    n_episodes_per_class=resolved_n_episodes,
                    force_reload=force_reload,
                    seed=seed,
                )
                dataset = safe_load(dset_loader, dataset_name, episode_size, current_task_name)
                if dataset is None:
                    continue
                actual_n_episodes = dset_loader.n_episodes_per_class
                cache_key = (dataset_name, episode_size, actual_n_episodes)
                if cache_key not in default_cache:
                    default_cache[cache_key] = extract_features(
                        dataset, episode_size, actual_n_episodes, batch_size, show_progress=progress_bar,
                    )
                current_X, current_y = default_cache[cache_key]

            # Resolve scores_path now that actual_n_episodes is known
            scores_path = os.path.join(
                output_folder, dset_str, model_str,
                f"{episode_size}_{actual_n_episodes}", current_task_name,
            )
            if tokens_suffix is not None:
                scores_path = os.path.join(scores_path, tokens_suffix)
            metrics_path = os.path.join(scores_path, "metrics.json")

            # Check for existing results after resolving "auto"
            if (
                resolved_n_episodes == "auto"
                and not force_rerun
                and not (force_rerun_oa and current_task_name == "order_alignment")
                and os.path.exists(metrics_path)
            ):
                print(colored(f"    -> Skipping (results already exist)", "yellow"))
                successes.append((dataset_name, episode_size, current_task_name))
                continue

            if "processor" in task_config:
                processor_module = importlib.import_module(f"steb.processors.{task_config['processor']}")
                processor_class_name = f"{task_config['processor'].replace('_', ' ').title().replace(' ', '')}Processor"
                processor_class = getattr(processor_module, processor_class_name)
                processor = processor_class()
            else:
                processor = Processor()

            processed_data = processor.process(current_X, current_y)

            task_module = importlib.import_module(f"steb.tasks.{current_task_name}")
            task_class_name = f"{current_task_name.replace('_', ' ').title().replace(' ', '')}Task"
            task_class = getattr(task_module, task_class_name)
            task = task_class()

            # LFTK uses abs-diff / L1-diff for pair tasks and clustering; others use cosine / K-Means
            from .models.lftk_model import LFTKModel

            is_lftk = isinstance(model, LFTKModel)
            if current_task_name in ("pre_defined_pair_classification", "all_to_all_pair_classification"):
                score_mode = "abs_diff" if is_lftk else "cosine"
                metrics = task.evaluate(*processed_data, score_mode=score_mode)
            elif current_task_name == "clustering":
                distance_mode = "l1_diff" if is_lftk else "euclidean"
                metrics = task.evaluate(*processed_data, distance_mode=distance_mode)
            else:
                metrics = task.evaluate(*processed_data)

            submetrics_config = task_config.get("submetrics", {})
            if submetrics_config:
                metrics["submetrics"] = _evaluate_submetrics(
                    submetrics_config,
                    processed_data,
                    task,
                )

            os.makedirs(scores_path, exist_ok=True)
            with open(metrics_path, "w+") as ouf:
                ouf.write(json.dumps(metrics))

            print(colored(f"    -> Metrics: {metrics}", "green"))
            successes.append((dataset_name, episode_size, current_task_name))

        except Exception as e:
            error_msg = f"{type(e).__name__}: {e}"
            print(colored(f"  FAILED {dataset_name}/{current_task_name}: {error_msg}", "red"))
            traceback.print_exc()
            failures.append((dataset_name, episode_size, current_task_name, error_msg))

    _print_evaluation_summary(successes, failures)

    model_str = os.path.basename(model.model_name_or_path)
    if model_str == "":
        model_str = os.path.basename(os.path.dirname(model.model_name_or_path))

    log_path = _write_evaluation_log(
        output_folder,
        model_str,
        run_name,
        successes,
        failures,
    )
    return {"successes": successes, "failures": failures, "log_path": log_path}

steb.core.preview(datasets, task_name=None, episode_sizes=None, n_episodes_per_class=None, output_file=None, show_summary=True)

Previews dataset statistics for a benchmark run without loading a model.

For each dataset/task/episode_size combination, reports class counts, resolved n_episodes_per_class, and which classes would be dropped.

Parameters:

Name	Type	Description	Default
datasets	List[str]	A list of dataset names to preview.	required
task_name	Optional[str]	The task to preview. If None, previews all tasks.	None
episode_sizes	Optional[List[int]]	Episode sizes to preview. If None, uses per-task defaults.	None
n_episodes_per_class	Optional[int]	Override for n_episodes_per_class. If None, uses per-task defaults.	None
output_file	Optional[str]	Optional path to write the preview report to.	None
show_summary	bool	Whether to print the summary block. Set to False when the caller handles its own summary (e.g. preset mode).	True

Returns:

Type	Description
List[Dict[str, Any]]	A list of result dicts, one per dataset/task/episode_size combination.

Source code in steb/core.py

def preview(
    datasets: List[str],
    task_name: Optional[str] = None,
    episode_sizes: Optional[List[int]] = None,
    n_episodes_per_class: Optional[int] = None,
    output_file: Optional[str] = None,
    show_summary: bool = True,
) -> List[Dict[str, Any]]:
    """
    Previews dataset statistics for a benchmark run without loading a model.

    For each dataset/task/episode_size combination, reports class counts,
    resolved n_episodes_per_class, and which classes would be dropped.

    Args:
        datasets: A list of dataset names to preview.
        task_name: The task to preview. If None, previews all tasks.
        episode_sizes: Episode sizes to preview. If None, uses per-task defaults.
        n_episodes_per_class: Override for n_episodes_per_class. If None, uses per-task defaults.
        output_file: Optional path to write the preview report to.
        show_summary: Whether to print the summary block. Set to False when
            the caller handles its own summary (e.g. preset mode).

    Returns:
        A list of result dicts, one per dataset/task/episode_size combination.
    """
    results = []
    lines = []

    for dataset_name, current_task_name, task_config, episode_size, resolved_n_episodes in _iter_task_configs(
        datasets, task_name, episode_sizes, n_episodes_per_class,
    ):
        if current_task_name is None:
            error_msg = task_config  # sentinel: error string stored in task_config slot
            print(colored(f"  {dataset_name} | ERROR: {error_msg}", "red"))
            continue

        try:
            loader = DatasetLoader(
                dataset_name=dataset_name,
                episode_size=episode_size,
                n_episodes_per_class=resolved_n_episodes,
                task_name=current_task_name,
            )
            stats = loader.preview()

            result = {
                "dataset": dataset_name,
                "task": current_task_name,
                **stats,
            }
            results.append(result)

            dropped_count = stats["dropped_classes"]
            status_str = "OK" if dropped_count == 0 else f"{dropped_count} dropped"
            plain_line = (
                f"  {dataset_name} | {current_task_name} | "
                f"ep_size={episode_size} | "
                f"n_episodes={stats['n_episodes_per_class']} | "
                f"classes={stats['kept_classes']}/{stats['total_classes']} | "
                f"min_count={stats['min_class_count']} | "
                f"{status_str}"
            )
            color = "green" if dropped_count == 0 else "yellow"
            print(colored(plain_line, color))
            lines.append(plain_line)

            if stats["dropped_labels"]:
                for label, count in sorted(stats["dropped_labels"].items(), key=lambda x: x[1]):
                    drop_line = f"    dropped '{label}': {count}/{stats['samples_per_class']} samples"
                    print(colored(drop_line, "yellow"))
                    lines.append(drop_line)

        except Exception as e:
            msg = f"  FAILED {dataset_name}/{current_task_name}: {type(e).__name__}: {e}"
            print(colored(msg, "red"))
            lines.append(msg)

    # Summary
    total_combos = len(results)
    total_dropped = sum(r["dropped_classes"] for r in results)
    total_kept = sum(r["kept_classes"] for r in results)
    total_total = sum(r["total_classes"] for r in results)

    if show_summary:
        summary_lines = [
            "",
            "=" * 60,
            "Preview Summary",
            "=" * 60,
            f"  Combinations: {total_combos}",
            f"  Classes kept: {total_kept}/{total_total}",
            f"  Classes dropped: {total_dropped}",
            "=" * 60,
        ]
        for line in summary_lines:
            print(colored(line, "cyan"))
        lines.extend(summary_lines)

    if output_file:
        with open(output_file, "w") as f:
            f.write("\n".join(lines) + "\n")
        print(colored(f"  Report saved to: {output_file}", "cyan"))

    return results

Base Model

steb.models.base.STEBModel

Bases: ABC

Abstract base class for text embedding models.

Subclasses that tokenize input may set effective_max_tokens in their constructor to the resolved per-text token cap when the user has opted into truncation mode or specified an explicit max_tokens cap. None means default chunk-and-pool behavior with no tagging.

Source code in steb/models/base.py

class STEBModel(ABC):
    """
    Abstract base class for text embedding models.

    Subclasses that tokenize input may set ``effective_max_tokens`` in
    their constructor to the resolved per-text token cap when the user
    has opted into truncation mode or specified an explicit ``max_tokens``
    cap. ``None`` means default chunk-and-pool behavior with no tagging.
    """

    effective_max_tokens: Optional[int] = None

    @abstractmethod
    def embed_multiple(
        self,
        episodes: List[List[str]],
        batch_size: int,
        show_progress: bool = False,
    ) -> np.ndarray:
        """
        Embeds a list of episodes, where each episode is a list of texts.

        Args:
            episodes: A list of episodes to embed.
            batch_size: The batch size to use for embedding.
            show_progress: Whether to show a progress bar.

        Returns:
            A numpy array of embeddings.
        """
        pass

embed_multiple(episodes, batch_size, show_progress=False) abstractmethod

Embeds a list of episodes, where each episode is a list of texts.

Parameters:

Name	Type	Description	Default
episodes	List[List[str]]	A list of episodes to embed.	required
batch_size	int	The batch size to use for embedding.	required
show_progress	bool	Whether to show a progress bar.	False

Returns:

Type	Description
ndarray	A numpy array of embeddings.

Source code in steb/models/base.py

@abstractmethod
def embed_multiple(
    self,
    episodes: List[List[str]],
    batch_size: int,
    show_progress: bool = False,
) -> np.ndarray:
    """
    Embeds a list of episodes, where each episode is a list of texts.

    Args:
        episodes: A list of episodes to embed.
        batch_size: The batch size to use for embedding.
        show_progress: Whether to show a progress bar.

    Returns:
        A numpy array of embeddings.
    """
    pass