Handle merges stored as list instead of space-separated string

arcee-ai · Oct 5, 2024 · ff76c62 · ff76c62
1 parent 8522917
commit ff76c62
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/mergekit/tokenizer/build.py b/mergekit/tokenizer/build.py
@@ -90,7 +90,12 @@ def get_stripped_tokenizer(
             del tok_dict["model"]["vocab"][tok]
 
     def _keep_merge(m):
-        toks = m.split(" ")
+        if isinstance(m, str) and m.count(" ") == 1:
+            toks = m.split(" ")
+        elif isinstance(m, list):
+            toks = m
+        else:
+            raise RuntimeError(f"Unexpected merge format: {repr(m)} ({type(m)})")
         for tok in toks:
             if tok in unused_toks:
                 return False