diff --git a/hyperlora/data_utils.py b/hyperlora/data_utils.py index b5ba802..664c1d8 100644 --- a/hyperlora/data_utils.py +++ b/hyperlora/data_utils.py @@ -1,5 +1,5 @@ import pandas as pd -from typing import Literal +from typing import Literal, List, Dict, Any, Tuple, Iterator, Callable, Optional, Union from transformers import PreTrainedTokenizerBase from training_utils import TRAINING_TASK @@ -10,7 +10,20 @@ IGNORE_INDEX = -100 def get_sft_prompt_formatting_fn( sft_mode: TRAINING_TASK, tokenizer: PreTrainedTokenizerBase, -): +) -> Callable[[Dict[str, Any]], Dict[str, Any]]: + """ + Get a function that formats examples for supervised fine-tuning. + + Args: + sft_mode: The training task type + tokenizer: The tokenizer to use for chat template application + + Returns: + A function that takes a training example and returns formatted data + + Raises: + NotImplementedError: If sft_mode is not COMPLETION or tokenizer has no chat template + """ if sft_mode != TRAINING_TASK.COMPLETION: raise NotImplementedError( f"Training task {sft_mode} not supported. " @@ -56,7 +69,23 @@ def get_sft_prompt_formatting_fn( return f_intx -def convert_ctx_prompt_response_to_messages(example, add_ctx_to_chat=True): +def convert_ctx_prompt_response_to_messages( + example: Dict[str, Any], + add_ctx_to_chat: bool = True, +) -> Dict[str, Any]: + """ + Convert context/prompt/response format to chat messages format. + + Args: + example: Dictionary containing 'prompt' and 'response' keys + add_ctx_to_chat: Whether to prepend context to the user message + + Returns: + Dictionary with added 'messages' key containing chat format + + Raises: + ValueError: If 'prompt' or 'response' keys are missing + """ if "prompt" not in example or "response" not in example: raise ValueError(f"'prompt' and 'response' are required. Got: {example}") @@ -77,7 +106,16 @@ def convert_ctx_prompt_response_to_messages(example, add_ctx_to_chat=True): return example -def get_preprocessing_fn(ds_name): +def get_preprocessing_fn(ds_name: str) -> Callable[[Dict[str, Any]], Dict[str, Any]]: + """ + Get preprocessing function for a specific dataset. + + Args: + ds_name: Name of the dataset + + Returns: + A preprocessing function that takes and returns a dictionary + """ f = lambda x: x # if ds_name == "context_numbers": @@ -191,16 +229,28 @@ def get_preprocessing_fn(ds_name): # taken from https://github.com/huggingface/trl/issues/632#issuecomment-1972630547 -def get_assistant_start_end_indices(messages, conversation_text): +def get_assistant_start_end_indices( + messages: List[Dict[str, str]], + conversation_text: str, +) -> List[Tuple[int, int]]: """ - Get the start and end indices of the assistant's messages in the conversation text. + Get the start and end indices of assistant messages in conversation text. + + Args: + messages: List of message dictionaries with 'role' and 'content' keys + conversation_text: Full conversation text + + Returns: + List of (start, end) index tuples for assistant messages """ start_indices = [] + # end_indices = [] current_index = 0 for message in messages: message_text = message["content"] match_index = conversation_text[current_index:].find(message_text) start_indices.append(current_index + match_index) + # end_indices.append(current_index + match_index + len(message_text)) current_index += match_index + len(message_text) end_indices = [ len(conversation_text) if i == len(start_indices) - 1 else start_indices[i + 1] @@ -212,7 +262,24 @@ def get_assistant_start_end_indices(messages, conversation_text): ] -def get_masked_labels(conversation_ids, assistant_ranges): +def get_masked_labels( + conversation_ids: Dict[str, List[Any]], + assistant_ranges: List[Tuple[int, int]], +) -> Iterator[int]: + """ + Generate masked labels for conversation, masking non-assistant tokens. + NOTE: This will also includes extra tokens between assistant and user messages + in multi-turn conversations. + E.g., {assistant_msg} <|eot_id|><|start_header_id|>user<|end_header_id|> {user_msg} + for Llama 3 models. + + Args: + conversation_ids: Dictionary with tokenized conversation info + assistant_ranges: List of (start, end) indices for assistant messages + + Yields: + Token ID or IGNORE_INDEX for each position + """ for id_, (id_s, id_e) in list( zip(conversation_ids["input_ids"], conversation_ids["offset_mapping"]) ): @@ -223,17 +290,30 @@ def get_masked_labels(conversation_ids, assistant_ranges): def tokenize_chat_messages( - example, - tokenizer, - mask_assistant_inputs=True, - tokenizer_kwargs=None, -): + example: Dict[str, Any], + tokenizer: PreTrainedTokenizerBase, + mask_assistant_inputs: bool = True, + tokenizer_kwargs: Optional[Dict[str, Any]] = None, +) -> Dict[str, List[int]]: + """ + Tokenize chat messages and optionally mask non-assistant tokens. + + Args: + example: Dictionary containing 'chat' and 'messages' keys + tokenizer: Tokenizer to use + mask_assistant_inputs: Whether to mask non-assistant tokens + tokenizer_kwargs: Additional arguments to pass to tokenizer + + Returns: + Dictionary containing tokenized inputs and labels + """ # should be used only with chat models text = example["chat"] messages = example["messages"] conversation_ids = tokenizer( text, return_offsets_mapping=mask_assistant_inputs, + add_special_tokens=False, **tokenizer_kwargs if tokenizer_kwargs is not None else {}, ) if mask_assistant_inputs: diff --git a/hyperlora/intx_sft.py b/hyperlora/intx_sft.py index f3f9815..c49fa29 100644 --- a/hyperlora/intx_sft.py +++ b/hyperlora/intx_sft.py @@ -8,6 +8,7 @@ from transformers import ( HfArgumentParser, TrainingArguments, DataCollatorForSeq2Seq, + EvalPrediction, ) @@ -21,7 +22,7 @@ from data_utils import ( from training_utils import TRAINING_TASK, train_model -def compute_metrics(eval_pred) -> dict: +def compute_metrics(eval_pred: EvalPrediction) -> dict: """ Custom metrics function for the trainer Args: @@ -137,7 +138,6 @@ def main(): "tokenizer": tokenizer, "mask_assistant_inputs": True, "tokenizer_kwargs": { - "add_special_tokens": False, "max_length": max_seq_len, }, },