use sharegpt from prompt strategies now

axolotl-ai-cloud · Sep 25, 2023 · 700ef36 · 700ef36
1 parent 7922664
commit 700ef36
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 13 deletions.
diff --git a/...lotl/prompt_strategies/sharegpt_simple.py → src/axolotl/prompt_strategies/sharegpt.py b/...lotl/prompt_strategies/sharegpt_simple.py → src/axolotl/prompt_strategies/sharegpt.py
@@ -1,12 +1,16 @@
 """Module containing the SimpleShareGPTPromptTokenizingStrategy class"""
+from typing import Any, Dict, Optional
 
 from axolotl.prompt_tokenizers import ShareGPTPromptTokenizingStrategy
 from axolotl.prompters import ShareGPTPrompterV2
 
 
-def load(tokenizer, cfg):
+def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
+    conversation = (
+        ds_cfg["conversation"] if ds_cfg and "conversation" in ds_cfg else None
+    )
     return SimpleShareGPTPromptTokenizingStrategy(
-        ShareGPTPrompterV2(),
+        ShareGPTPrompterV2(conversation=conversation),
         tokenizer,
         cfg.train_on_inputs,
         cfg.sequence_len,

diff --git a/src/axolotl/utils/data.py b/src/axolotl/utils/data.py
@@ -25,7 +25,6 @@
     GPTeacherPromptTokenizingStrategy,
     JeopardyPromptTokenizingStrategy,
     OpenAssistantPromptTokenizingStrategy,
-    ShareGPTPromptTokenizingStrategy,
     SummarizeTLDRPromptTokenizingStrategy,
 )
 from axolotl.prompters import (
@@ -35,7 +34,6 @@
     MultipleChoiceConcisePrompter,
     MultipleChoiceExplainPrompter,
     ReflectAlpacaPrompter,
-    ShareGPTPrompterV2,
     SummarizeTLDRPrompter,
 )
 from axolotl.utils.dict import DictDefault
@@ -320,15 +318,6 @@ def for_d_in_datasets(dataset_configs):
                 )
                 ds_wrapper = TokenizedPromptDataset(ds_strategy, ds)
                 datasets.append(ds_wrapper)
-            elif d_base_type == "sharegpt":
-                ds_strategy = ShareGPTPromptTokenizingStrategy(
-                    ShareGPTPrompterV2(d.conversation),
-                    tokenizer,
-                    cfg.train_on_inputs,
-                    cfg.sequence_len,
-                )
-                ds_wrapper = TokenizedPromptDataset(ds_strategy, ds)
-                datasets.append(ds_wrapper)
             else:
                 suffix = ""
                 if ":load_" in d.type: