| from termcolor import colored | |
| import logging | |
| def check_dataset_labels(dataset, tokenizer): | |
| # the dataset is already shuffled, so let's just check the first 5 elements | |
| for idx in range(5): | |
| check_example_labels(dataset[idx], tokenizer) | |
| def check_example_labels(example, tokenizer): | |
| # Get the input_ids, labels, and attention_mask from the dataset | |
| input_ids = example["input_ids"] | |
| labels = example["labels"] | |
| attention_mask =example["attention_mask"] | |
| # You can compare the input_ids and labels element-wise | |
| # Remember to ignore positions with IGNORE_TOKEN_ID (if you use it) or attention_mask equal to 0 | |
| colored_tokens = [] | |
| for i, (input_id, label_id, mask) in enumerate( | |
| zip(input_ids, labels, attention_mask) | |
| ): | |
| decoded_input_token = tokenizer.decode(input_id) | |
| # Choose the color based on whether the label has the ignore value or not | |
| color = ( | |
| "red" if label_id == -100 else ("yellow" if label_id == 0 else "green") | |
| ) | |
| colored_token = colored(decoded_input_token, color) + colored( | |
| f"({label_id}, {mask}, {input_id})", "white" | |
| ) | |
| colored_tokens.append(colored_token) | |
| logging.info(" ".join(colored_tokens)) | |
| logging.info("\n\n\n") | |