INNER CODE UNIT · Python
create_labels
CogStack/OpenGPT · opengpt/dataset_utils.py:173
def create_labels(examples, config, tokenizer):
r''' This is used with a prepared HF dataset that is already tokenized. It will add labels
so that only the AI generated parts (answers) will be trained on.
'''
user_token_id = tokenizer.vocab[config.special_tokens.user]
ai_token_id = tokenizer.vocab[config.special_tokens.ai]
# Everything written by an AI will be used for training, and everything by a user will be ignored
examples['labels'] = []
for i in range(len(examples['input_ids'])):
labels = []
ignore = True
for tkn_id in examples['input_ids'][i]:
if tkn_id == user_token_id:
ignore = True
elif tkn_id == ai_token_id:
ignore = False