winglian commited on
Commit
ba043a3
·
unverified ·
1 Parent(s): 41353d2

add ultrachat prompt strategies (#996)

Browse files
src/axolotl/prompt_strategies/sharegpt.py CHANGED
@@ -39,6 +39,23 @@ def load(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
39
  return strategy
40
 
41
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
42
  def load_role(tokenizer, cfg):
43
  return SimpleRoleShareGPTPromptTokenizingStrategy(
44
  ShareGPTPrompterV2(),
@@ -109,3 +126,17 @@ class GuanacoShareGPTPromptTokenizingStrategy(ShareGPTPromptTokenizingStrategy):
109
  {"from": role_map[t["role"]], "value": t["text"]} for t in conversations
110
  ]
111
  return turns
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
39
  return strategy
40
 
41
 
42
+ def load_ultrachat(tokenizer, cfg, ds_cfg: Optional[Dict[str, Any]] = None):
43
+ conversation = (
44
+ ds_cfg["conversation"] if ds_cfg and "conversation" in ds_cfg else None
45
+ )
46
+ strategy = UltrachatShareGPTPromptTokenizingStrategy(
47
+ ShareGPTPrompterV2(
48
+ conversation=conversation,
49
+ ),
50
+ tokenizer,
51
+ cfg.train_on_inputs,
52
+ cfg.sequence_len,
53
+ )
54
+ if ds_cfg and "strict" in ds_cfg:
55
+ strategy.strict = ds_cfg["strict"]
56
+ return strategy
57
+
58
+
59
  def load_role(tokenizer, cfg):
60
  return SimpleRoleShareGPTPromptTokenizingStrategy(
61
  ShareGPTPrompterV2(),
 
126
  {"from": role_map[t["role"]], "value": t["text"]} for t in conversations
127
  ]
128
  return turns
129
+
130
+
131
+ class UltrachatShareGPTPromptTokenizingStrategy(SimpleShareGPTPromptTokenizingStrategy):
132
+ """
133
+ sharegpt strategy that remaps ultrachat data to sharegpt format
134
+ """
135
+
136
+ def get_conversation_thread(self, prompt):
137
+ conversations = prompt["messages"]
138
+ role_map = {"user": "human", "assistant": "gpt"}
139
+ turns = [
140
+ {"from": role_map[t["role"]], "value": t["content"]} for t in conversations
141
+ ]
142
+ return turns