feat: support template's stop_str as list (#2678)

lm-sys · Nov 13, 2023 · aeec0e0 · aeec0e0
1 parent a333a55
commit aeec0e0
Showing 1 changed file with 12 additions and 1 deletion.
diff --git a/fastchat/llm_judge/gen_model_answer.py b/fastchat/llm_judge/gen_model_answer.py
@@ -141,8 +141,19 @@ def get_model_answers(
                         output_ids,
                         spaces_between_special_tokens=False,
                     )
-                    if conv.stop_str and output.find(conv.stop_str) > 0:
+                    if conv.stop_str and isinstance(conv.stop_str, list):
+                        stop_str_indices = sorted(
+                            [
+                                output.find(stop_str)
+                                for stop_str in conv.stop_str
+                                if output.find(stop_str) > 0
+                            ]
+                        )
+                        if len(stop_str_indices) > 0:
+                            output = output[: stop_str_indices[0]]
+                    elif conv.stop_str and output.find(conv.stop_str) > 0:
                         output = output[: output.find(conv.stop_str)]
+
                     for special_token in tokenizer.special_tokens_map.values():
                         if isinstance(special_token, list):
                             for special_tok in special_token: