Preprocess conversation data for the model input. Parameters: sources (List[Dict]): A list of conversation segments. tokenizer (PreTrainedTokenizer): A tokenizer instance. max_len (int): The maximum sequence length. system_message (str, optional): A default
(
sources,
tokenizer: transformers.PreTrainedTokenizer,
max_len: int,
system_message: str = "You are a helpful assistant."
)
| 28 | IGNORE_TOKEN_ID = LabelSmoother.ignore_index |
| 29 | |
| 30 | def preprocess( |
| 31 | sources, |
| 32 | tokenizer: transformers.PreTrainedTokenizer, |
| 33 | max_len: int, |
| 34 | system_message: str = "You are a helpful assistant." |
| 35 | ) -> Dict: |
| 36 | """ |
| 37 | Preprocess conversation data for the model input. |
| 38 | |
| 39 | Parameters: |
| 40 | sources (List[Dict]): A list of conversation segments. |
| 41 | tokenizer (PreTrainedTokenizer): A tokenizer instance. |
| 42 | max_len (int): The maximum sequence length. |
| 43 | system_message (str, optional): A default system message. |
| 44 | |
| 45 | Returns: |
| 46 | Dict: A dictionary with 'input_ids', 'labels', and 'attention_mask'. |
| 47 | """ |
| 48 | roles = {"user": "<|im_start|>user", "assistant": "<|im_start|>assistant"} |
| 49 | im_start = tokenizer.im_start_id |
| 50 | im_end = tokenizer.im_end_id |
| 51 | nl_tokens = tokenizer('\n').input_ids |
| 52 | _system = tokenizer('system').input_ids + nl_tokens |
| 53 | _user = tokenizer('user').input_ids + nl_tokens |
| 54 | _assistant = tokenizer('assistant').input_ids + nl_tokens |
| 55 | |
| 56 | # Apply prompt templates |
| 57 | input_ids, targets = [], [] |
| 58 | for i, source in enumerate(sources): |
| 59 | if roles[source[0]["from"]] != roles["user"]: |
| 60 | source = source[1:] |
| 61 | |
| 62 | input_id, target = [], [] |
| 63 | system = [im_start] + _system + tokenizer(system_message).input_ids + [im_end] + nl_tokens |
| 64 | input_id += system |
| 65 | target += [im_start] + [IGNORE_TOKEN_ID] * (len(system)-3) + [im_end] + nl_tokens |
| 66 | assert len(input_id) == len(target) |
| 67 | for sentence in enumerate(source): |
| 68 | role = roles[sentence["from"]] |
| 69 | _input_id = tokenizer(role).input_ids + nl_tokens + \ |
| 70 | tokenizer(sentence["value"]).input_ids + [im_end] + nl_tokens |
| 71 | input_id += _input_id |
| 72 | if role == '<|im_start|>user': |
| 73 | _target = [im_start] + [IGNORE_TOKEN_ID] * (len(_input_id)-3) + [im_end] + nl_tokens |
| 74 | elif role == '<|im_start|>assistant': |
| 75 | _target = [im_start] + [IGNORE_TOKEN_ID] * len(tokenizer(role).input_ids) + \ |
| 76 | _input_id[len(tokenizer(role).input_ids)+1:-2] + [im_end] + nl_tokens |
| 77 | else: |
| 78 | raise NotImplementedError |
| 79 | target += _target |
| 80 | assert len(input_id) == len(target) |
| 81 | input_id += [tokenizer.pad_token_id] * (max_len - len(input_id)) |
| 82 | target += [IGNORE_TOKEN_ID] * (max_len - len(target)) |
| 83 | input_ids.append(input_id[:max_len]) |
| 84 | targets.append(target[:max_len]) |
| 85 | input_ids = torch.tensor(input_ids, dtype=torch.int) |
| 86 | targets = torch.tensor(targets, dtype=torch.int) |
| 87 |