Fix missing sys handling

instructlab · Jun 20, 2024 · cc556f7 · cc556f7
1 parent c27e20b
commit cc556f7
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/instructlab/training/data_process.py b/src/instructlab/training/data_process.py
@@ -196,7 +196,10 @@ def main(args: DataProcessArgs):
 
     eos_tk = get_sp_token(tokenizer, SPECIAL_TOKENS.eos)
     pad_tk = get_sp_token(tokenizer, SPECIAL_TOKENS.pad)
-    system_tk = get_sp_token(tokenizer, SPECIAL_TOKENS.system)
+    if SPECIAL_TOKENS.system:
+        system_tk = get_sp_token(tokenizer, SPECIAL_TOKENS.system)
+    else:
+        system_tk = None
     user_tk = get_sp_token(tokenizer, SPECIAL_TOKENS.user)
     assistant_tk = get_sp_token(tokenizer, SPECIAL_TOKENS.assistant)
     log_rank_0(