integrations.kd.chat_template

integrations.kd.chat_template

Chat template prompt strategy loader with KD support

Classes

Name	Description
ChatTemplateStrategyWithKD	Handle fields for logprob KD
ChatTemplateStrategyWithKDv2	Strat for datasets with complete structured KD logprob data
KDStrategyLoader	Load ChatTemplateStrategy with KD support using StrategyLoader.
KDStrategyLoaderV2	Load KD chat template datasets with pre-tokenized logprob data

ChatTemplateStrategyWithKD

integrations.kd.chat_template.ChatTemplateStrategyWithKD(
    prompter,
    tokenizer,
    train_on_inputs,
    sequence_len,
    roles_to_train=None,
    train_on_eos=None,
    train_on_eot=None,
    eot_tokens=None,
    split_thinking=False,
    logprobs_field='logprobs',
    gen_temperature=1.0,
    kd_temperature=1.0,
)

Handle fields for logprob KD

Methods

Name	Description
transform_logprobs	Transform logprobs to target format for KD training

transform_logprobs

integrations.kd.chat_template.ChatTemplateStrategyWithKD.transform_logprobs(
    sample,
)

Transform logprobs to target format for KD training

ChatTemplateStrategyWithKDv2

integrations.kd.chat_template.ChatTemplateStrategyWithKDv2(
    prompter,
    tokenizer,
    train_on_inputs,
    sequence_len,
    roles_to_train=None,
    train_on_eos=None,
    train_on_eot=None,
    eot_tokens=None,
    split_thinking=False,
    logprobs_field='logprobs',
    gen_temperature=1.0,
    kd_temperature=1.0,
)

Strat for datasets with complete structured KD logprob data

Methods

Name	Description
transform_logprobs	Transform logprobs to target format for KD training

transform_logprobs

integrations.kd.chat_template.ChatTemplateStrategyWithKDv2.transform_logprobs(
    sample,
)

Transform logprobs to target format for KD training

KDStrategyLoader

integrations.kd.chat_template.KDStrategyLoader()

Load ChatTemplateStrategy with KD support using StrategyLoader.

KDStrategyLoaderV2

integrations.kd.chat_template.KDStrategyLoaderV2()

Load KD chat template datasets with pre-tokenized logprob data