Upload slimshard.py

slimshard.py

@@ -1546,6 +1546,15 @@ def _write_lm_head_calibrated_chunked(
             seqs_to_run.append(chat_seq[:-1])
             start_pos = len(chat_pref_ids) - 1
             seq_transition_slices.append((start_pos, chat_rep_ids))
+        elif tok_id("<｜User｜>") >= 0 and len(chat_rep_ids) >= 2:
+            ds_pref = [0, 128799] + tok_piece_ids(
+                ["Reasoning", " Effort", ":", " ", "7", "5", " (", "range", " ", "1", "-", "1", "0", "0", ",", " the", " higher", " the", " value", ",", " the", " more", " thorough", " the", " reasoning", ")\n\n"]
+            ) + [tok_id("<｜User｜>")] + tok_piece_ids(["hello"]) + [tok_id("<｜Assistant｜>"), tok_id("<think>")]
+            ds_rep = ([tok_id("</think>")] if tok_id("</think>") >= 0 else []) + chat_rep_ids
+            if all(x >= 0 for x in ds_pref):
+                ds_seq = ds_pref + ds_rep
+                seqs_to_run.append(ds_seq[:-1])
+                seq_transition_slices.append((len(ds_pref) - 1, ds_rep))
 
         bg_seq = [int(x) for x in np.linspace(100, min(max(vocab - 1, 100), 20000), 16, dtype=np.int64)]
         seqs_to_run.append(bg_seq)