Implementation:Hpcaitech ColossalAI Eval BaseModel

Knowledge Sources	Hpcaitech_ColossalAI
Domains	Evaluation, Benchmarking
Last Updated	2026-02-09 00:00 GMT

Overview

BaseModel is the abstract base class for all model wrappers in the ColossalEval evaluation framework, defining the interface for inference, generation, and loss computation.

Description

The class provides the foundational structure that all ColossalEval model wrappers must implement. It initializes with a model path, maximum sequence length (default 2048), a prompt template (defaulting to the "plain" conversation template), batch size, and a logger. The class declares three abstract methods: inference for running full inference on data (calling both generate and model forward), generate for producing text completions given input strings, and get_loss for computing loss at target tokens by masking based on tokenization length differences. A to convenience method moves the underlying model to a specified device.

Usage

Use BaseModel as the parent class when creating a new model wrapper for the ColossalEval framework. Concrete implementations should override inference, generate, and get_loss to provide model-specific behavior.

Code Reference

Source Location

Repository: Hpcaitech_ColossalAI
File: applications/ColossalEval/colossal_eval/models/base.py
Lines: 1-79

Signature

class BaseModel:
    def __init__(
        self,
        path: str,
        model_max_length: int = 2048,
        prompt_template: Conversation = None,
        batch_size: int = 1,
        logger: DistributedLogger = None,
    ):

    @abstractclassmethod
    def inference(self, data: List[Dict]) -> None:

    @abstractclassmethod
    def generate(self, inputs: List[str], max_new_tokens: int) -> List[str]:

    @abstractclassmethod
    def get_loss(self, batch: List[str], batch_target: List[str]) -> List[float]:

    def to(self, device):

Import

from colossal_eval.models.base import BaseModel

I/O Contract

Inputs (init)

Name	Type	Required	Description
path	str	Yes	Path to the model weights or model identifier
model_max_length	int	No	Maximum sequence length the model supports (default 2048)
prompt_template	Conversation	No	Conversation template for prompt formatting (defaults to "plain" template)
batch_size	int	No	Batch size for inference (default 1)
logger	DistributedLogger	No	Logger instance for distributed logging

Outputs

Name	Type	Description
inference return	None	Modifies data in-place with model outputs
generate return	List[str]	List of generated text strings
get_loss return	List[float]	List of loss values for each sample in the batch

Usage Examples

from colossal_eval.models.base import BaseModel
from colossal_eval.utils import prompt_templates

# BaseModel is abstract; use a concrete subclass
# Example of how a subclass would be initialized:
# model = MyConcreteModel(
#     path="/path/to/model",
#     model_max_length=4096,
#     prompt_template=prompt_templates["alpaca"],
#     batch_size=8,
# )
# model.inference(data)

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment