diff --git a/configs/base.yaml b/configs/base.yaml index fae2e4c..f0359b6 100644 --- a/configs/base.yaml +++ b/configs/base.yaml @@ -73,6 +73,7 @@ inference: top_p: 0.1 temperature: 0.6 model: ${model.checkpoint} + num_display_generations: 5 evaluation: additional_metrics: diff --git a/configs/data/aeslc.yaml b/configs/data/aeslc.yaml index b69281a..2b15b50 100644 --- a/configs/data/aeslc.yaml +++ b/configs/data/aeslc.yaml @@ -12,3 +12,4 @@ is_in_conversational_format: false system_prompt: "Write a subject text for the email.\n\n" assistant_response_start: ${model.assistant_response_start} use_test_benchmark: false +task: 'summarization' diff --git a/configs/data/example.yaml b/configs/data/example.yaml index 7f14ee6..18bc134 100644 --- a/configs/data/example.yaml +++ b/configs/data/example.yaml @@ -5,7 +5,7 @@ few_shot: count: 0 unlabeled_data_split_name: str test_split_name: str -input_column_name: str +input_column_name: str or dict[str, str] input_max_length: int output_column_name: str or list[str] output_max_length: int @@ -15,3 +15,4 @@ system_prompt: str assistant_response_start: str is_in_conversational_format: bool use_test_benchmark: bool +task: 'str' diff --git a/configs/data/gigaword.yaml b/configs/data/gigaword.yaml index bcca145..311813c 100644 --- a/configs/data/gigaword.yaml +++ b/configs/data/gigaword.yaml @@ -10,3 +10,4 @@ output_max_length: 21 fetch_kwargs: {} system_prompt: "Generate a headline for the article in lowercase." use_test_benchmark: false +task: 'summarization' diff --git a/configs/data/gsm8k.yaml b/configs/data/gsm8k.yaml new file mode 100644 index 0000000..0e95e73 --- /dev/null +++ b/configs/data/gsm8k.yaml @@ -0,0 +1,15 @@ +dataset: ['openai/gsm8k', 'main'] +input_column_name: 'question' +output_column_name: 'answer' +unlabeled_data_split_name: train +test_split_name: test +train_subset_size: null +test_subset_size: null +input_max_length: 2048 +output_max_length: 1024 +fetch_kwargs: {} +is_in_conversational_format: false +system_prompt: "Please generate a solution to the math problem below. In the end, add a line break and print the answer with \"#### \" replacing the tag with the answer." +assistant_response_start: ${model.assistant_response_start} +use_test_benchmark: false +task: 'math' diff --git a/configs/data/math.yaml b/configs/data/math.yaml new file mode 100644 index 0000000..5242970 --- /dev/null +++ b/configs/data/math.yaml @@ -0,0 +1,15 @@ +dataset: 'Aktsvigun/preprocessed_for_al_OpenR1-Math-220k' +input_column_name: 'problem' +output_column_name: {'train': 'short_solution', 'test': 'answer'} +unlabeled_data_split_name: train +test_split_name: test +train_subset_size: null +test_subset_size: 1_000 +input_max_length: 4096 +output_max_length: 4096 +fetch_kwargs: {} +is_in_conversational_format: false +system_prompt: "Please solve the math problem. In the end, add two line breaks and print the answer with \"Answer: \"." +assistant_response_start: ${model.assistant_response_start} +use_test_benchmark: false +task: 'math' diff --git a/configs/data/race.yaml b/configs/data/race.yaml new file mode 100644 index 0000000..e53b527 --- /dev/null +++ b/configs/data/race.yaml @@ -0,0 +1,27 @@ +dataset: ['ehovy/race', 'all'] +input_column_name: {"article": "article", "question": "question", "options": "options"} +output_column_name: 'answer' +unlabeled_data_split_name: train +test_split_name: test +train_subset_size: null +test_subset_size: 1_000 +input_max_length: 1024 +output_max_length: 1 +fetch_kwargs: {} +is_in_conversational_format: false +system_prompt: Please answer the question based on the article. Answer only with the letter corresponding to the correct answer. Only output the letter of the correct answer and nothing else. +user_prompt_template: | + Article: + {article} + + Question: + {question} + + Options: + {options} + + Correct answer (only the letter): + +assistant_response_start: ${model.assistant_response_start} +use_test_benchmark: false +task: 'multi-choice-qa' diff --git a/configs/data/test.yaml b/configs/data/test.yaml index b87f481..fab7c9e 100644 --- a/configs/data/test.yaml +++ b/configs/data/test.yaml @@ -14,3 +14,4 @@ train_subset_size: 100 is_in_conversational_format: false assistant_response_start: ${model.assistant_response_start} use_test_benchmark: false +task: 'summarization' diff --git a/configs/data/trivia_qa.yaml b/configs/data/trivia_qa.yaml index aa09f15..94ab6d4 100644 --- a/configs/data/trivia_qa.yaml +++ b/configs/data/trivia_qa.yaml @@ -12,3 +12,4 @@ is_in_conversational_format: false system_prompt: "Please answer the question very laconically:\n\n" assistant_response_start: ${model.assistant_response_start} use_test_benchmark: false +task: 'open-qa' diff --git a/configs/data/user_data.yaml b/configs/data/user_data.yaml index e18be97..ee2d098 100644 --- a/configs/data/user_data.yaml +++ b/configs/data/user_data.yaml @@ -11,3 +11,4 @@ system_prompt: "Write a summary for the article in lowercase." test_subset_size: 100 train_subset_size: 100 use_test_benchmark: false +task: 'summarization' diff --git a/configs/data/xsum.yaml b/configs/data/xsum.yaml index ea93a93..64ea51b 100644 --- a/configs/data/xsum.yaml +++ b/configs/data/xsum.yaml @@ -12,3 +12,4 @@ fetch_kwargs: is_in_conversational_format: false system_prompt: "Summarize the article in one sentence." use_test_benchmark: false +task: 'summarization' diff --git a/configs/labeller/api_llm.yaml b/configs/labeller/api_llm.yaml index 1217908..11803b5 100644 --- a/configs/labeller/api_llm.yaml +++ b/configs/labeller/api_llm.yaml @@ -1,6 +1,7 @@ type: api_llm provider: openai api_key: 'api_key.key' +base_url: null mode: sync # or batched parameters: model: gpt-4o diff --git a/configs/test.yaml b/configs/test.yaml index e762094..da93b8d 100644 --- a/configs/test.yaml +++ b/configs/test.yaml @@ -62,6 +62,7 @@ inference: temperature: 0.6 top_p: 0.1 model: ${model.checkpoint} + num_display_generations: 0 evaluation: additional_metrics: [] diff --git a/src/atgen/labellers/api_labellers/anthropic_labeller.py b/src/atgen/labellers/api_labellers/anthropic_labeller.py index 648a86e..bda71a5 100644 --- a/src/atgen/labellers/api_labellers/anthropic_labeller.py +++ b/src/atgen/labellers/api_labellers/anthropic_labeller.py @@ -31,9 +31,9 @@ def __call__(self, dataset: Dataset) -> Dataset: annotations = [] price = 0 - for text in tqdm(data): + for messages in tqdm(data): request_kwargs = dict(base_request_kwargs) - request_kwargs["messages"] = [{"role": "user", "content": text}] + request_kwargs["messages"] = messages output = self.client.messages.create(**request_kwargs) price += self._calculate_price(output) annotations.append(output.content[0].text.strip()) diff --git a/src/atgen/labellers/api_labellers/openai_labeller.py b/src/atgen/labellers/api_labellers/openai_labeller.py index 0c0530e..8ebef95 100644 --- a/src/atgen/labellers/api_labellers/openai_labeller.py +++ b/src/atgen/labellers/api_labellers/openai_labeller.py @@ -22,38 +22,6 @@ MAX_NUM_TRIES = 3 UPDATE_TIME_IN_SECONDS = 10 # update time when checking for the completion - -# SYSTEM_PROMPT = """ -# Act as an experienced ... -# [INSERT] -# Examples for few-shot learning are below. -# """.strip() - -# EXAMPLE_1_INPUT = """ -# [INSERT] -# """.lstrip() -# -# EXAMPLE_2_INPUT = """ -# [INSERT] -# """.lstrip() -# -# EXAMPLE_1_OUTPUT = """ -# [INSERT] -# """.strip() -# -# EXAMPLE_2_OUTPUT = """ -# [INSERT] -# """.strip() - -# Can remove / add examples. If you remove them completely, do not forget to remove the last line from the system prompt -messages_template = [ - # {"role": "system", "content": SYSTEM_PROMPT}, - # {"role": "system", "name": "example_user", "content": EXAMPLE_1_INPUT}, - # {"role": "system", "name": "example_assistant", "content": EXAMPLE_1_OUTPUT}, - {"role": "user", "content": ""}, -] - - class OpenAILabeller(BaseLabeler): def __init__( self, @@ -65,7 +33,10 @@ def __init__( super().__init__(output_column_name, budget) self.config = config # Create the OpenAI client - kwargs = {} if base_url is None else {"base_url": base_url} + if base_url := (base_url or config.base_url): + kwargs = {"base_url": base_url} + else: + kwargs = {} self.client = OpenAI(api_key=self.config.api_key, **kwargs) self.mode = config.get("mode") @@ -75,16 +46,12 @@ def _sync_call(self, dataset: Dataset) -> Dataset: annotations = [] total_price = 0 - for text in tqdm(data, desc="Processing with OpenAI API"): - # Prepare messages for this input - text_messages = deepcopy(messages_template) - text_messages[-1]["content"] = text - + for messages in tqdm(data, desc="Processing with OpenAI API"): # Make API call with retries for attempt in range(MAX_NUM_TRIES): try: response = self.client.chat.completions.create( - messages=text_messages, **self.config.parameters + messages=messages, **self.config.parameters ) break except Exception as e: @@ -135,6 +102,10 @@ def _sync_call(self, dataset: Dataset) -> Dataset: dataset = dataset.remove_columns(self.output_column_name) dataset = dataset.add_column(self.output_column_name, annotations) + # TODO: make a parameter + for i, annotation in enumerate(annotations[:2]): + print(f"Annotation of the {i+1}th instance: {annotation}") + print("-"*100) return dataset def _batched_call(self, dataset: Dataset) -> Dataset: diff --git a/src/atgen/metrics/compute_metrics.py b/src/atgen/metrics/compute_metrics.py index f55c175..2864372 100644 --- a/src/atgen/metrics/compute_metrics.py +++ b/src/atgen/metrics/compute_metrics.py @@ -1,6 +1,9 @@ +import string from time import time import logging +from typing import Literal, Optional from omegaconf import DictConfig +import re import numpy as np from evaluate import load @@ -24,6 +27,7 @@ def compute_metrics( generated_texts, reference_texts, original_texts, + task: Literal["summarization", "open-qa", "multi-choice-qa", "translation", "math"], config: DictConfig, cache_dir: str = "cache", ) -> dict[str, float]: @@ -34,6 +38,7 @@ def compute_metrics( generated_texts: List of generated texts to evaluate reference_texts: List of reference texts (ground truth) or list of lists of reference texts original_texts: List of source texts + task: Task type (summarization, open-qa, multi-choice-qa, translation) config: Configuration for evaluation - additional_metrics: List of additional metrics to use. Options include: - "bartscore": BARTScore metrics @@ -71,25 +76,35 @@ def compute_metrics( metric = AnswerRelevanceMetric(model=llm) ``` """ - # Load metrics that are always used - sacrebleu = load("sacrebleu", cache_dir=cache_dir) - rouge = load("rouge", cache_dir=cache_dir) + if task == "multi-choice-qa": + metrics_to_calculate = ["exact_match"] + list(config.additional_metrics) + elif task == "open-qa": + metrics_to_calculate = ["exact_match"] + list(config.additional_metrics) + elif task == "summarization": + metrics_to_calculate = ["exact_match", "sacrebleu", "bleu", "rouge", "word_length"] + list(config.additional_metrics) + elif task == "translation": + metrics_to_calculate = ["exact_match", "sacrebleu", "bleu", "word_length"] + list(config.additional_metrics) + elif task == "math": + metrics_to_calculate = ["exact_match_math"] + list(config.additional_metrics) + else: + raise NotImplementedError(f"Task {task} not implemented") + + if "sacrebleu" in metrics_to_calculate: + sacrebleu = load("sacrebleu", cache_dir=cache_dir) + if "rouge" in metrics_to_calculate: + rouge = load("rouge", cache_dir=cache_dir) result = {} - result["word_length_gen"] = np.array( - [len(text.split()) for text in generated_texts] - ) + if "word_length" in metrics_to_calculate: + result["word_length_gen"] = np.array( + [len(text.split()) for text in generated_texts] + ) time_dict = {} # Metrics that use both the generated texts and the original texts and # those that do not require reference texts - src_word_lengths = np.array([len(text.split()) for text in original_texts]) - - # Avoid division by zero - src_word_lengths_safe = np.where(src_word_lengths > 0, src_word_lengths, 1) - result["word_length_src_rel"] = result["word_length_gen"] / src_word_lengths_safe - if "bartscore" in config.additional_metrics and is_bart_score_available: + if "bartscore" in metrics_to_calculate and is_bart_score_available: log.info("Calculating BARTScore scores...") start_time = time() result.update( @@ -105,70 +120,88 @@ def compute_metrics( # Metrics that use both the generated texts and the reference texts if reference_texts is not None: # Exact match - if isinstance(reference_texts[0], list): - result["exact_match"] = np.array( + if "exact_match" in metrics_to_calculate: + if isinstance(reference_texts[0], list): + result["exact_match"] = np.array( + [ + any(_preprocess_text(pred) == _preprocess_text(one_ref) for one_ref in ref) + for pred, ref in zip(generated_texts, reference_texts) + ] + ) + else: + result["exact_match"] = np.array( + [_preprocess_text(pred) == _preprocess_text(ref) for pred, ref in zip(generated_texts, reference_texts)] + ) + if "exact_match_math" in metrics_to_calculate: + # result["exact_match_math"] = np.array( + # [ + # pred.split("Answer: ")[-1].lower() == ref.lower() + # for pred, ref in zip(generated_texts, reference_texts) + # ] + # ) + result["exact_match_math"] = np.array( [ - any(pred == one_ref for one_ref in ref) + pred.split("#### ")[-1].lower() == ref.split("#### ")[-1].lower() for pred, ref in zip(generated_texts, reference_texts) ] ) - else: - result["exact_match"] = np.array( - [pred == ref for pred, ref in zip(generated_texts, reference_texts)] - ) - # BLEU - start_time = time() - result["bleu"] = np.array( - [ - pair_bleu(references=ref, prediction=pred) - for pred, ref in tqdm(zip(generated_texts, reference_texts)) - ] - ) - time_dict["time_bleu"] = time() - start_time - # ROUGE - start_time = time() - result.update( - rouge.compute( - predictions=generated_texts, - references=reference_texts, - use_stemmer=True, + if "bleu" in metrics_to_calculate: + # BLEU + start_time = time() + result["bleu"] = np.array( + [ + pair_bleu(references=ref, prediction=pred) + for pred, ref in tqdm(zip(generated_texts, reference_texts)) + ] ) - ) - time_dict["time_rouge"] = time() - start_time - # Sacrebleu - start_time = time() - if not isinstance(reference_texts[0], list): - sacrebleu_references = [[ref] for ref in reference_texts] - sacrebleu_result = sacrebleu.compute( - predictions=generated_texts, references=sacrebleu_references + time_dict["time_bleu"] = time() - start_time + if "rouge" in metrics_to_calculate: + # ROUGE + start_time = time() + result.update( + rouge.compute( + predictions=generated_texts, + references=reference_texts, + use_stemmer=True, + ) ) - result["sacrebleu"] = sacrebleu_result.pop("score") - else: - sacrebleu_scores = [] - for pred, ref in zip(generated_texts, reference_texts): + time_dict["time_rouge"] = time() - start_time + if "sacrebleu" in metrics_to_calculate: + # Sacrebleu + start_time = time() + if not isinstance(reference_texts[0], list): + sacrebleu_references = [[ref] for ref in reference_texts] sacrebleu_result = sacrebleu.compute( - predictions=[pred], references=[ref] + predictions=generated_texts, references=sacrebleu_references ) - sacrebleu_scores.append(sacrebleu_result.pop("score")) - result["sacrebleu"] = sacrebleu_scores + result["sacrebleu"] = sacrebleu_result.pop("score") + else: + sacrebleu_scores = [] + for pred, ref in zip(generated_texts, reference_texts): + sacrebleu_result = sacrebleu.compute( + predictions=[pred], references=[ref] + ) + sacrebleu_scores.append(sacrebleu_result.pop("score")) + result["sacrebleu"] = sacrebleu_scores - time_dict["time_sacrebleu"] = time() - start_time - # Lengths - if isinstance(reference_texts[0], list): - ref_word_lengths = np.array( - [ - np.mean([len(text.split()) for text in ref]) - for ref in reference_texts - ] - ) - else: - ref_word_lengths = np.array([len(ref.split()) for ref in reference_texts]) - # Avoid division by zero - ref_word_lengths_safe = np.where(ref_word_lengths > 0, ref_word_lengths, 1) - result["word_length_rel"] = result["word_length_gen"] / ref_word_lengths_safe + time_dict["time_sacrebleu"] = time() - start_time + if "word_length" in metrics_to_calculate: + # Lengths + if isinstance(reference_texts[0], list): + ref_word_lengths = np.array( + [ + np.mean([len(text.split()) for text in ref]) + for ref in reference_texts + ] + ) + else: + ref_word_lengths = np.array([len(ref.split()) for ref in reference_texts]) + # Avoid division by zero + ref_word_lengths_safe = np.where(ref_word_lengths > 0, ref_word_lengths, 1) + result["word_length_rel"] = result["word_length_gen"] / ref_word_lengths_safe # AlignScore - if "alignscore" in config.additional_metrics and is_alignscore_available: + if "alignscore" in metrics_to_calculate and is_alignscore_available: log.info("Calculating AlignScore scores...") start_time = time() alignscores = calculate_alignscore( @@ -241,3 +274,32 @@ def compute_metrics( } return result + +def _preprocess_text(text: str, do_lowercase: bool = True, do_remove_punctuation: bool = True, do_remove_extra_spaces: bool = True, do_remove_stopwords: bool = False, stopwords: Optional[list[str]] = None) -> str: + # Convert to lowercase + if do_lowercase: + text = text.lower() + + # Remove punctuation + if do_remove_punctuation: + # Keep hyphens within words, remove other punctuation + text = re.sub(r'(? Dataset: + alphabet_titled = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" + input_column_names = data_config.input_column_name + options_column_name = input_column_names["options"] + user_prompt_template = data_config.user_prompt_template + system_prompt = data_config.system_prompt + messages = [] + for inst in dataset: + preprocessed_options = '' + for option, letter in zip(inst[options_column_name], alphabet_titled): + preprocessed_options += f"- {letter}. {option}\n" + user_prompt_kwargs = { + key: inst[key] for key in input_column_names.keys() if key != options_column_name + } + user_prompt_kwargs["options"] = preprocessed_options + inst_messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": user_prompt_template.format(**user_prompt_kwargs)}, + ] + if split == "train": + inst_messages.append({"role": "assistant", "content": inst["answer"]}) + messages.append(inst_messages) + if "messages" in dataset.column_names: + dataset = dataset.remove_columns(["messages"]) + dataset = dataset.add_column("messages", messages) + # Can't directly update `processed_input_column_name` because test data is loaded separately + OmegaConf.update(data_config, "processed_input_column_name", "messages", force_add=True) + data_config.is_in_conversational_format = True + return dataset diff --git a/src/atgen/utils/data/prepare_conversational_data.py b/src/atgen/utils/data/prepare_conversational_data.py index 432d9c3..2bc07f5 100644 --- a/src/atgen/utils/data/prepare_conversational_data.py +++ b/src/atgen/utils/data/prepare_conversational_data.py @@ -1,4 +1,3 @@ -from itertools import chain from datasets import Dataset from omegaconf import DictConfig diff --git a/src/atgen/utils/generate.py b/src/atgen/utils/generate.py index ef080c2..7bd7638 100644 --- a/src/atgen/utils/generate.py +++ b/src/atgen/utils/generate.py @@ -115,6 +115,7 @@ def generate_vllm( del llm_runner gc.collect() cuda.empty_cache() + _maybe_display_generations(generations, inference_config.num_display_generations) return generations @@ -213,6 +214,7 @@ def generate_response(s, messages): ) # Clean up engine.shutdown() + _maybe_display_generations(generations, inference_config.num_display_generations) return generations @@ -280,6 +282,7 @@ def generate_transformers( model_name=model.name_or_path, framework=TRANSFORMERS_FRAMEWORK ) + _maybe_display_generations(generations, inference_config.num_display_generations) return generations @@ -339,3 +342,9 @@ def tokenize_conversational_example( input_ids = tokenizer.apply_chat_template(example["messages"], add_generation_prompt=True) attention_mask = [1 for _ in range(len(input_ids))] return {"input_ids": input_ids, "attention_mask": attention_mask} + +def _maybe_display_generations(generations: list[str], num_display_gens: int): + if num_display_gens: + log.info("Displaying the first {} generations:".format(num_display_gens)) + for i in range(num_display_gens): + log.info(f"Generation {i+1}: {generations[i]}") diff --git a/src/atgen/utils/resolvers.py b/src/atgen/utils/resolvers.py index b123220..4d4959e 100644 --- a/src/atgen/utils/resolvers.py +++ b/src/atgen/utils/resolvers.py @@ -2,9 +2,7 @@ Custom Hydra resolvers for configuration calculations """ -from typing import Union - -from omegaconf import OmegaConf, DictConfig, ListConfig +from omegaconf import OmegaConf def multiply_with_few_shot(input_max_length: int, few_shot_count: int) -> int: diff --git a/src/atgen/utils/training_utils.py b/src/atgen/utils/training_utils.py index 1646d3f..4816e71 100644 --- a/src/atgen/utils/training_utils.py +++ b/src/atgen/utils/training_utils.py @@ -100,9 +100,13 @@ def torch_call(self, examples): ) if not response_token_ids_idxs: - raise RuntimeError( - f"Could not find response key {self.response_token_ids} in the input_ids" + # Log warning and skip this example by masking everything + logger.warning( + f"Could not find response key {self.response_token_ids} in example {i}. " + f"Masking entire example to exclude from loss." ) + batch["labels"][i, :] = self.ignore_index + continue # If there's an instruction template, find those too human_token_ids_idxs = [] @@ -111,44 +115,57 @@ def torch_call(self, examples): batch["input_ids"][i], self.instruction_token_ids ) + # Ensure sequence ends with instruction_template + last assistant response + last_response_idx = response_token_ids_idxs[-1] + + # Find the user message before the last assistant response + preceding_human_idxs = [ + idx for idx in human_token_ids_idxs if idx < last_response_idx + ] + if not preceding_human_idxs: + # Log warning and skip this example by masking everything + logger.warning( + f"No user message found before the last assistant response in example {i}. " + f"Masking entire example to exclude from loss." + ) + batch["labels"][i, :] = self.ignore_index + continue + + # Find any content after the last assistant response + next_human_idxs = [ + idx for idx in human_token_ids_idxs if idx > last_response_idx + ] + + # If there's content after the last assistant response, truncate it + if next_human_idxs: + end_idx = next_human_idxs[0] + # Truncate the sequence + batch["input_ids"][i, end_idx:] = self.tokenizer.pad_token_id + batch["attention_mask"][i, end_idx:] = 0 + # Set all labels to ignore_index as default batch["labels"][i, :] = self.ignore_index # Only unmask the last assistant response - if len(response_token_ids_idxs) > 0: - last_response_idx = response_token_ids_idxs[-1] - - # Skip the template tokens - only unmask content after the template - content_start_idx = last_response_idx + len(self.response_token_ids) - - # Determine end of the response - if len(human_token_ids_idxs) > 0: - # Find the next user message after the last assistant message, if any - next_human_idxs = [ - idx for idx in human_token_ids_idxs if idx > last_response_idx - ] - if next_human_idxs: - end_idx = next_human_idxs[0] - else: - end_idx = batch["input_ids"].shape[1] # End of sequence - else: - end_idx = batch["input_ids"].shape[1] # End of sequence - - # Unmask only the content after the template and before the end - batch["labels"][i, content_start_idx:end_idx] = batch["input_ids"][ - i, content_start_idx:end_idx - ] + content_start_idx = last_response_idx + len(self.response_token_ids) + end_idx = batch["input_ids"].shape[1] + + # Find the actual end of the response (before padding) + if self.tokenizer.pad_token_id is not None: + padding_mask = batch["input_ids"][i] == self.tokenizer.pad_token_id + if padding_mask.any(): + end_idx = padding_mask.nonzero()[0].item() + + # Unmask only the content after the template and before the end + batch["labels"][i, content_start_idx:end_idx] = batch["input_ids"][ + i, content_start_idx:end_idx + ] # Always mask padding tokens if self.tokenizer.pad_token_id is not None: padding_mask = batch["input_ids"][i] == self.tokenizer.pad_token_id batch["labels"][i, padding_mask] = self.ignore_index - # TODO: probably remove this - # # Also use attention mask to identify padding (if available) - # if "attention_mask" in batch: - # batch["labels"][i, batch["attention_mask"][i] == 0] = self.ignore_index - return batch def _find_token_ids_indices(self, token_ids, pattern): @@ -258,7 +275,20 @@ def _dataset_to_chat_template( tokenized_texts = _formatting_fn(dataset, tokenizer) texts = [] for text in tokenized_texts: - if not data_collator.response_template in text: + # Check if the tokenized response template is present + # This ensures we only keep examples that will work with the collator + tokenized = tokenizer(text, truncation=True, return_tensors="pt") + input_ids = tokenized["input_ids"][0].tolist() + + # Check if response_token_ids exist in the tokenized input + response_token_ids = data_collator.response_token_ids + found = False + for i in range(len(input_ids) - len(response_token_ids) + 1): + if input_ids[i : i + len(response_token_ids)] == response_token_ids: + found = True + break + + if not found: text = "" texts.append(text) return {TEXT_FIELD: texts}