Kkit.llm_utils.data_process

View Source

 1def input_output_to_messages(input_key, output_key, dataset, drop_input_output=True):
 2    def _fn(examples):
 3        return {
 4            "messages": [
 5                {"role": "user", "content": examples[input_key]},
 6                {"role": "assistant", "content": examples[output_key]}
 7            ]
 8        }
 9    
10    return dataset.map(_fn, remove_columns=[input_key, output_key] if drop_input_output else None)

def input_output_to_messages(input_key, output_key, dataset, drop_input_output=True): View Source

 2def input_output_to_messages(input_key, output_key, dataset, drop_input_output=True):
 3    def _fn(examples):
 4        return {
 5            "messages": [
 6                {"role": "user", "content": examples[input_key]},
 7                {"role": "assistant", "content": examples[output_key]}
 8            ]
 9        }
10    
11    return dataset.map(_fn, remove_columns=[input_key, output_key] if drop_input_output else None)