tocmo0nlord/axolotl
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

tweak check for batched prompt data

Browse Source
This commit is contained in:
Wing Lian
2025-01-07 14:54:32 -05:00
parent 684b38291f
commit 6784822cfb
2 changed files with 9 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
src/axolotl/prompt_strategies/bradley_terry/chat_template.py
View File

@@ -21,6 +21,9 @@ class BTChatTemplateStrategy(ChatTemplateStrategy):
Bradley-Terry reward model pairwise chat template prompt strategy. Bradley-Terry reward model pairwise chat template prompt strategy.
""" """
def is_prompt_batched(self, prompt: dict[str, Any]) -> bool:
return all(isinstance(v, list) for v in prompt.values())
def _tokenize_single_prompt(self, prompt): def _tokenize_single_prompt(self, prompt):
""" """

7
src/axolotl/prompt_strategies/chat_template.py
View File

@@ -226,12 +226,17 @@ class ChatTemplateStrategy(PromptTokenizingStrategy):
# Let calling code know we can handle lists of examples # Let calling code know we can handle lists of examples
return True return True
def is_prompt_batched(self, prompt: dict[str, Any]) -> bool:
return all(isinstance(v, list) for v in prompt.values()) and all(
isinstance(v, list) for v in prompt[self.messages]
)
def tokenize_prompt(self, prompt: dict[str, Any]): def tokenize_prompt(self, prompt: dict[str, Any]):
""" """
Public method that can handle either a single prompt or a batch of prompts. Public method that can handle either a single prompt or a batch of prompts.
""" """
if not all(isinstance(v, list) for v in prompt.values()): if not self.is_prompt_batched(prompt):
return self._tokenize_single_prompt(prompt) return self._tokenize_single_prompt(prompt)
res = defaultdict(lambda: []) res = defaultdict(lambda: [])