From 13f31d837a84600ed38919614230d459b2e53769 Mon Sep 17 00:00:00 2001 From: NanoCode012 Date: Thu, 18 Jan 2024 18:07:52 +0900 Subject: [PATCH] fix: add handling of empty role back --- src/axolotl/prompt_strategies/sharegpt.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/axolotl/prompt_strategies/sharegpt.py b/src/axolotl/prompt_strategies/sharegpt.py index aaa199661d..fa7e6ae10c 100644 --- a/src/axolotl/prompt_strategies/sharegpt.py +++ b/src/axolotl/prompt_strategies/sharegpt.py @@ -277,6 +277,17 @@ def tokenize_prompt(self, prompt): labels[:len_role] = [IGNORE_TOKEN_ID] * min( len_role, len(labels) ) + elif role == "": + turn = content + # 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 + ) + if self.train_on_inputs: + labels = copy.deepcopy(res["input_ids"]) + else: + # everything from this is masked out from the labels + labels = [IGNORE_TOKEN_ID] * len(res["input_ids"]) else: LOG.warning(f"unhandled role: {role}") continue