Implementation:Hpcaitech ColossalAI Eval Conversation
| Knowledge Sources | |
|---|---|
| Domains | Evaluation, Benchmarking |
| Last Updated | 2026-02-09 00:00 GMT |
Overview
Conversation is a dataclass-based prompt template system that manages conversation formatting for different model architectures in the ColossalEval framework, supporting multiple separator styles.
Description
The module defines a SeparatorStyle enum with four formatting styles: ADD_BOS_EOS_TOKEN (wraps messages with and tokens), ALPACA (uses "### Instruction" / "### Response" role prefixes), PLAIN (no formatting, concatenates messages directly), and YAYI (uses <|Human|> / <|YaYi|> role prefixes). The Conversation dataclass provides methods for prompt construction (get_prompt), prompt-with-target generation (get_prompt_with_target), message management (append_message, clear), and serialization (copy, dict, save_prompt). The module also provides two key utility functions: get_few_shot_prefix which constructs a few-shot prompt prefix that fits within a token budget, and get_batch_prompt which processes batches of data samples into formatted prompts, handling both single-turn and multi-turn conversations. Four pre-built conversation templates are exported: conv_coati, conv_alpaca, conv_plain, and conv_yayi.
Usage
Use Conversation objects and the prompt_templates dictionary to format prompts for different model architectures during ColossalEval evaluation. The get_batch_prompt function is the primary entry point for converting evaluation data into model-ready prompts.
Code Reference
Source Location
- Repository: Hpcaitech_ColossalAI
- File: applications/ColossalEval/colossal_eval/utils/conversation.py
- Lines: 1-265
Signature
class SeparatorStyle(Enum):
ADD_BOS_EOS_TOKEN = auto()
ALPACA = auto()
PLAIN = auto()
YAYI = auto()
@dataclasses.dataclass
class Conversation:
system: str
roles: List[str]
messages: List[List[str]]
offset: int
sep_style: SeparatorStyle = SeparatorStyle.ADD_BOS_EOS_TOKEN
sep: str = "</s>"
def clear(self):
def get_prompt(self):
def get_prompt_with_target(self, target):
def save_prompt(self):
def append_message(self, role, message):
def copy(self):
def dict(self):
def get_few_shot_prefix(few_shot_data: List[str], tokenizer: Optional[AutoTokenizer], max_tokens: int) -> str:
def get_batch_prompt(conv: Conversation, batch: List[Dict], few_shot_data: List[str], tokenizer: Optional[AutoTokenizer], model_max_length: Optional[int]) -> Tuple[List[Dict], List[Dict]]:
prompt_templates = {"coati": conv_coati, "alpaca": conv_alpaca, "plain": conv_plain, "yayi": conv_yayi}
Import
from colossal_eval.utils.conversation import Conversation, SeparatorStyle, get_batch_prompt, prompt_templates
I/O Contract
Inputs (get_batch_prompt)
| Name | Type | Required | Description |
|---|---|---|---|
| conv | Conversation | Yes | Conversation template instance used for prompt formatting |
| batch | List[Dict] | Yes | List of data sample dictionaries with fields instruction, input, output, target, and dataset |
| few_shot_data | List[str] | No | List of few-shot example strings; first element is a description, rest are examples |
| tokenizer | Optional[AutoTokenizer] | No | Tokenizer for computing token lengths (required when few_shot_data is provided) |
| model_max_length | Optional[int] | No | Maximum model sequence length for few-shot token budget calculation |
Outputs (get_batch_prompt)
| Name | Type | Description |
|---|---|---|
| batch_prompt | List[str] | List of formatted prompt strings ready for model input |
| batch_target | List[List[str]] | List of target answer lists (each target is a list of possible answers) |
Usage Examples
from colossal_eval.utils.conversation import Conversation, SeparatorStyle, prompt_templates, get_batch_prompt
# Use a pre-built template
conv = prompt_templates["alpaca"].copy()
conv.append_message(conv.roles[0], "What is 2+2?")
conv.append_message(conv.roles[1], None)
print(conv.get_prompt())
# Process a batch of data
batch = [{"instruction": "Solve this:", "input": "2+2=?", "output": "", "target": "4", "dataset": "math"}]
prompts, targets = get_batch_prompt(conv, batch, few_shot_data=None, tokenizer=None, model_max_length=None)