File size: 16,377 Bytes
5d86137
 
ce24f5e
5e37144
1d5ab84
 
be22551
ce24f5e
e7d3e2d
1eebbd0
ce24f5e
e7d3e2d
 
 
5e37144
 
553a86b
 
ce24f5e
cb9797e
a1f9850
 
 
ce24f5e
e7d3e2d
 
ce24f5e
8d959a7
5d86137
 
 
8d959a7
 
ce24f5e
5d86137
 
 
 
ce24f5e
 
 
 
 
 
 
 
 
 
 
97d3776
ce24f5e
 
 
 
 
97d3776
 
 
 
5d86137
1d5ab84
2a428e8
 
 
 
 
 
1d5ab84
 
5d86137
1d5ab84
2a428e8
 
 
 
 
 
1d5ab84
 
1eebbd0
 
 
 
8662e8f
1eebbd0
 
 
 
 
 
97d3776
1eebbd0
 
 
85cf4f8
 
8e46c0f
85cf4f8
 
97d3776
8e46c0f
 
 
 
 
9d629d8
 
 
 
 
8e46c0f
 
 
 
 
 
ce24f5e
87d7825
5d86137
 
 
 
8d20e0a
 
 
87d7825
 
ce24f5e
5d86137
 
 
 
 
7925ddc
 
 
 
 
ce34d64
 
7925ddc
 
 
 
ce24f5e
7925ddc
 
 
 
 
 
 
ce24f5e
7925ddc
ce24f5e
5d86137
 
 
ce34d64
 
 
 
 
 
 
 
 
ce24f5e
 
87d7825
5d86137
 
 
 
 
87d7825
ce24f5e
87d7825
 
 
 
 
b46bc02
5d86137
 
 
 
 
b46bc02
 
 
1365073
b46bc02
 
 
a12fb0a
5d86137
 
 
 
 
a12fb0a
 
 
 
 
 
 
87d7825
5d86137
 
 
 
 
87d7825
 
 
 
 
 
 
1365073
5d86137
 
 
 
 
1365073
 
 
 
 
 
 
87d7825
5d86137
 
 
 
 
87d7825
 
 
ce24f5e
 
 
 
6045345
5d86137
 
 
 
 
6045345
 
 
 
 
 
 
81de0ef
5d86137
 
 
 
 
81de0ef
 
 
2bc1a5b
 
5d86137
2bc1a5b
 
 
 
 
 
 
81de0ef
 
ce34d64
 
 
 
 
 
 
 
81de0ef
 
 
 
 
 
 
 
 
5d86137
 
 
ce34d64
 
 
 
 
 
 
 
 
 
 
81de0ef
8e46c0f
81de0ef
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
5d86137
 
 
 
 
81de0ef
 
 
 
 
 
 
 
6045345
ce24f5e
5d86137
 
 
 
21c8e2d
 
 
ce24f5e
8e46c0f
1d5ab84
 
e7d3e2d
 
 
8d959a7
5d86137
21c8e2d
ce34d64
1d5ab84
e7d3e2d
a363604
1d5ab84
a363604
 
ce34d64
e7d3e2d
37293dc
 
ce34d64
1d5ab84
4ea9a66
1d5ab84
ce34d64
e7d3e2d
4ea9a66
a363604
 
 
 
ce34d64
e7d3e2d
37293dc
 
ce34d64
1d5ab84
37293dc
 
 
 
1d5ab84
 
e7d3e2d
 
25eeeeb
 
e7d3e2d
25eeeeb
 
 
1d5ab84
553a86b
e7d3e2d
8e46c0f
 
 
 
 
 
 
 
 
5e37144
5d86137
 
5e37144
 
a363604
 
 
 
 
 
 
 
 
 
 
5e37144
a363604
 
5e37144
 
 
 
 
 
a363604
 
 
 
 
5e37144
 
 
 
 
8e46c0f
 
 
 
 
 
 
e9650d3
8e46c0f
 
 
 
 
 
 
 
 
 
 
 
77bdb7d
be22551
8e46c0f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
"""Module containing PromptTokenizingStrategy and Prompter classes"""

import abc
import copy
import functools
import logging
from typing import Dict, List, Tuple, Union

from fastchat.conversation import Conversation
from transformers import BatchEncoding, PreTrainedTokenizer

from axolotl.monkeypatch.fastchat_conversation_turns import (
    add_get_turns_to_conversation,
)
from axolotl.prompters import IGNORE_TOKEN_ID

LOG = logging.getLogger("axolotl")

IGNORE_INDEX = -100
LLAMA_DEFAULT_PAD_TOKEN = "<pad>"  # nosec
LLAMA_DEFAULT_EOS_TOKEN = "</s>"  # nosec
LLAMA_DEFAULT_BOS_TOKEN = "<s>"  # nosec
LLAMA_DEFAULT_UNK_TOKEN = "<unk>"  # nosec

add_get_turns_to_conversation()


class InvalidDataException(Exception):
    """
    Exception raised when the data is invalid
    """


class PromptTokenizingStrategy(abc.ABC):
    """
    Abstract class for tokenizing strategies
    """

    def __init__(
        self,
        prompter,
        tokenizer,
        train_on_inputs: bool = False,
        sequence_len: int = 2048,
    ):
        self.prompter = prompter
        self.tokenizer: PreTrainedTokenizer = tokenizer
        self.train_on_inputs = train_on_inputs
        self.sequence_len = sequence_len
        self.max_length = sequence_len

    @abc.abstractmethod
    def tokenize_prompt(self, prompt):
        pass

    @property
    def supports_batched(self):
        return False

    @functools.lru_cache(maxsize=128)
    def _get_user_token(self):
        try:
            id_or_ids = self.tokenizer.convert_tokens_to_ids("<|USER|>")
            if isinstance(id_or_ids, (int,)):
                return id_or_ids
        except KeyError:
            pass
        return False

    @functools.lru_cache(maxsize=128)
    def _get_assistant_token(self):
        try:
            id_or_ids = self.tokenizer.convert_tokens_to_ids("<|ASSISTANT|>")
            if isinstance(id_or_ids, (int,)):
                return id_or_ids
        except KeyError:
            pass
        return False

    def _tokenize(
        self, prompt: str, add_eos_token: bool = True, strip_bos_token: bool = False
    ) -> BatchEncoding:
        result: BatchEncoding
        if not prompt:
            LOG.warning("Empty text requested for tokenization.")
            result = BatchEncoding(data={"input_ids": [], "attention_mask": []})
        else:
            result = self.tokenizer(
                prompt,
                truncation=True,
                max_length=self.max_length,
                padding=False,
                return_tensors=None,
            )
        if len(result["input_ids"]) == 0:
            LOG.warning("Tokenizer result is empty. You may want to audit your dataset")
        if (
            len(result["input_ids"]) > 0
            and result["input_ids"][-1] != self.tokenizer.eos_token_id
            and len(result["input_ids"]) < self.max_length
            and add_eos_token
        ):
            result["input_ids"].append(self.tokenizer.eos_token_id)
            result["attention_mask"].append(1)

        if (
            len(result["input_ids"]) > 0
            and result["input_ids"][0] == self.tokenizer.bos_token_id
            and strip_bos_token
        ):
            result["input_ids"] = result["input_ids"][1:]
            result["attention_mask"] = result["attention_mask"][1:]

        result["labels"] = result["input_ids"].copy()
        return result


class InstructionPromptTokenizingStrategy(PromptTokenizingStrategy):
    """
    Tokenizing strategy for instruction-based prompts.
    """

    def parse_instruction_fields(
        self, prompt
    ) -> Union[Tuple[str, str, str], Tuple[str, str, str, str]]:
        raise NotImplementedError

    def tokenize_prompt(self, prompt):
        (
            instruction,
            input,  # pylint: disable=redefined-builtin
            response,
        ) = self.parse_instruction_fields(prompt)
        user_prompt = next(
            iter(
                self.prompter.build_prompt(
                    instruction,
                    input,
                )
            )
        )
        tokenized_prompt = self._tokenize(user_prompt, add_eos_token=False)
        if not self.train_on_inputs:
            user_prompt_len = len(tokenized_prompt["input_ids"])
            # TODO this could be sped up using numpy array slicing
            tokenized_prompt["labels"] = [-100] * user_prompt_len
        tokenized_res_prompt = self._tokenize(
            response, strip_bos_token=True, add_eos_token=True
        )
        tokenized_prompt["input_ids"] += tokenized_res_prompt["input_ids"]
        tokenized_prompt["attention_mask"] += tokenized_res_prompt["attention_mask"]
        tokenized_prompt["labels"] += tokenized_res_prompt["input_ids"]

        return tokenized_prompt

    def _build_full_prompt(
        self, instruction, input, response  # pylint: disable=redefined-builtin
    ):
        return next(
            iter(
                self.prompter.build_prompt(
                    instruction,
                    input,
                    response,
                )
            )
        )


class AlpacaPromptTokenizingStrategy(InstructionPromptTokenizingStrategy):
    """
    Tokenizing strategy for Alpaca prompts.
    """

    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str]:
        return (
            prompt["instruction"],
            prompt["input"] if "input" in prompt else "",
            prompt["output"],
        )


class AlpacaMultipleChoicePromptTokenizingStrategy(InstructionPromptTokenizingStrategy):
    """
    Tokenizing strategy for Alpaca Multiple Choice prompts.
    """

    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str]:
        return (
            prompt["question"],
            "\n".join(f'- "{choice}"' for choice in prompt["choices"]),
            prompt["solution"] if "solution" in prompt else prompt["explanation"],
        )


class JeopardyPromptTokenizingStrategy(InstructionPromptTokenizingStrategy):
    """
    Tokenizing strategy for Jeopardy prompts.
    """

    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str]:
        return (
            prompt["question"],
            prompt["category"],
            "what is " + prompt["answer"],
        )


class OpenAssistantPromptTokenizingStrategy(InstructionPromptTokenizingStrategy):
    """
    Tokenizing strategy for OpenAssistant prompts.
    """

    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str]:
        return (
            prompt["INSTRUCTION"],
            "",
            prompt["RESPONSE"],
        )


class SummarizeTLDRPromptTokenizingStrategy(InstructionPromptTokenizingStrategy):
    """
    Tokenizing strategy for SummarizeTLDR prompts.
    """

    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str]:
        return (
            prompt["article"],
            "",
            prompt["summary"],
        )


class GPTeacherPromptTokenizingStrategy(InstructionPromptTokenizingStrategy):
    """
    Tokenizing strategy for GPTeacher prompts.
    """

    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str]:
        return (
            prompt["instruction"],
            prompt["input"] if "input" in prompt else "",
            prompt["response"],
        )


class NomicGPT4AllPromptTokenizingStrategy(InstructionPromptTokenizingStrategy):
    """
    Tokenizing strategy for NomicGPT4All prompts.
    """

    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str]:
        return (
            prompt["prompt"],
            "",
            prompt["response"],
        )


class ReflectionPromptTokenizingStrategy(PromptTokenizingStrategy):
    """
    Tokenizing strategy for Reflection prompts.
    """

    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str, str, str]:
        raise NotImplementedError

    def tokenize_prompt(self, prompt):
        (
            instruction,
            input,  # pylint: disable=redefined-builtin
            output,
            reflection,
            corrected,
        ) = self.parse_instruction_fields(prompt)
        full_prompt = self._build_full_prompt(
            instruction, input, output, reflection, corrected
        )
        tokenized_full_prompt = self._tokenize(full_prompt)
        if not self.train_on_inputs:
            user_prompt = next(
                iter(
                    self.prompter.build_prompt(
                        instruction,
                        input,
                    )
                )
            )
            tokenized_user_prompt = self._tokenize(user_prompt, add_eos_token=False)
            user_prompt_len = len(tokenized_user_prompt["input_ids"])
            # TODO this could be sped up using numpy array slicing
            tokenized_full_prompt["labels"] = [
                -100
            ] * user_prompt_len + tokenized_full_prompt["labels"][user_prompt_len:]

        return tokenized_full_prompt

    def _build_full_prompt(
        self, instruction, input, output, reflection, corrected
    ):  # pylint: disable=redefined-builtin
        return next(
            iter(
                self.prompter.build_prompt(
                    instruction,
                    input,
                    output,
                    reflection,
                    corrected,
                )
            )
        )

    def _tokenize(self, prompt, add_eos_token=True, strip_bos_token=False):
        result = self.tokenizer(
            prompt,
            truncation=True,
            max_length=self.sequence_len,
            padding=False,
            return_tensors=None,
        )
        if (
            result["input_ids"][-1] != self.tokenizer.eos_token_id
            and len(result["input_ids"]) < self.sequence_len
            and add_eos_token
        ):
            result["input_ids"].append(self.tokenizer.eos_token_id)
            result["attention_mask"].append(1)

        result["labels"] = result["input_ids"].copy()
        return result


class AlpacaReflectionPTStrategy(ReflectionPromptTokenizingStrategy):
    """
    Tokenizing strategy for Alpaca Reflection prompts.
    """

    def parse_instruction_fields(self, prompt) -> Tuple[str, str, str, str, str]:
        return (
            prompt["instruction"],
            prompt["input"] if "input" in prompt else "",
            prompt["output"],
            prompt["reflection"],
            prompt["corrected"],
        )


class ShareGPTPromptTokenizingStrategy(PromptTokenizingStrategy):
    """
    Tokenizing strategy for ShareGPT prompts.
    """

    def get_conversation_thread(self, prompt):
        return prompt["conversations"]

    def tokenize_prompt(self, prompt):
        result, current_len = tokenize_prompt_default()
        user_token = self._get_user_token()
        assistant_token = self._get_assistant_token()
        conversation: Conversation = (
            self.prompter._conversation  # pylint: disable=protected-access
        )
        try:
            for _, part in enumerate(
                self.prompter.build_prompt(self.get_conversation_thread(prompt))
            ):
                if isinstance(part, tuple):
                    if conversation.roles[0] in part[0]:
                        turn = part[0] + part[1] if not user_token else part[1]
                        # this is still the user query, we should
                        if not part[1].strip():
                            LOG.warning(f"user turn has empty text: {prompt}")
                        res = self._tokenize(
                            turn,
                            add_eos_token=False,
                            strip_bos_token=True,
                        )
                        if user_token:
                            res["input_ids"] = [user_token, *res["input_ids"]]
                        # everything from this is masked out from the labels
                        labels = [IGNORE_TOKEN_ID] * len(res["input_ids"])
                    elif conversation.roles[1] in part[0]:
                        # TODO label assistant token/tokens w/ IGNORE_TOKEN_ID
                        turn = part[0] + part[1] if not assistant_token else part[1]
                        # this should be the assistant response, should end with an eos token
                        if not part[1].strip():
                            LOG.warning(f"assistant turn has empty text: {prompt}")
                        res = self._tokenize(
                            turn,
                            add_eos_token=True,
                            strip_bos_token=True,
                        )
                        if assistant_token:
                            res["input_ids"] = [
                                assistant_token,
                                *res["input_ids"],
                            ]
                        # not masked out from labels
                        labels = copy.deepcopy(res["input_ids"])
                    elif part[0] == "":
                        turn = part[1]
                        # this is only ever the first part, should include the bos token and the user query
                        res = self._tokenize(
                            turn, add_eos_token=False, strip_bos_token=False
                        )
                        # everything from this is masked out from the labels
                        labels = [IGNORE_TOKEN_ID] * len(res["input_ids"])
                    else:
                        LOG.warning(f"unhandled role: {part[0]}")
                        continue

                # pylint: disable=duplicate-code
                result, current_len = parse_tokenized_to_result(
                    result,
                    current_len,
                    res,
                    labels,
                    pad_token_id=self.tokenizer.pad_token_id,
                )
            return result
        except (KeyError, AssertionError, IndexError) as err:
            raise InvalidDataException(str(err)) from err

    def _tokenize(self, prompt, add_eos_token=True, strip_bos_token=False):
        if not prompt.strip():
            LOG.warning("Empty text requested for tokenization.")
            result = BatchEncoding(data={"input_ids": [], "attention_mask": []})
        else:
            result = self.tokenizer(
                prompt,
                truncation=True,
                max_length=self.sequence_len,
                padding=False,
                return_tensors=None,
            )
        if (
            len(result["input_ids"]) > 0
            and result["input_ids"][-1] != self.tokenizer.eos_token_id
            and len(result["input_ids"]) < self.sequence_len
            and add_eos_token
        ):
            result["input_ids"].append(self.tokenizer.eos_token_id)
            result["attention_mask"].append(1)

        if (
            len(result["input_ids"]) > 0
            and result["input_ids"][0] == self.tokenizer.bos_token_id
            and strip_bos_token
        ):
            result["input_ids"] = result["input_ids"][1:]
            result["attention_mask"] = result["attention_mask"][1:]

        result["labels"] = result["input_ids"].copy()
        return result


def tokenize_prompt_default() -> Tuple[Dict[str, List[int]], int]:
    """
    Returns the default values for the tokenize prompt function
    """

    result: Dict[str, List[int]] = {
        "input_ids": [],
        "attention_mask": [],
        "labels": [],
    }
    current_len = 0
    return result, current_len


def parse_tokenized_to_result(
    result: Dict[str, List[int]],
    current_len: int,
    res: Dict[str, List[int]],
    labels: List[int],
    pad_token_id: Union[int, None] = None,
) -> Tuple[Dict[str, List[int]], int]:
    """
    Parses the tokenized prompt and append the tokenized input_ids, attention_mask and labels to the result
    """

    input_ids = res["input_ids"]
    input_len = len(input_ids)
    result["input_ids"][current_len : current_len + input_len] = input_ids
    result["attention_mask"][current_len : current_len + input_len] = [
        1 if x != pad_token_id else 0 for x in input_ids
    ]
    result["labels"][current_len : current_len + input_len] = labels
    current_len += input_len

    return result, current_len