diff --git a/modeling/stoppers.py b/modeling/stoppers.py index 94c09e85..3f277f48 100644 --- a/modeling/stoppers.py +++ b/modeling/stoppers.py @@ -126,7 +126,12 @@ class Stoppers: # null_character = model.tokenizer.encode(chr(0))[0] if "completed" not in model.gen_state: model.gen_state["completed"] = [False] * len(input_ids) - + if utils.koboldai_vars.adventure: + extra_options = ["> You", "You:", "\n\n You", "\n\nYou", ". You"] + for option in extra_options: + if option not in utils.koboldai_vars.stop_sequence: + utils.koboldai_vars.stop_sequence.append(option) + #one issue is that the stop sequence may not actual align with the end of token #if its a subsection of a longer token for stopper in utils.koboldai_vars.stop_sequence: @@ -140,6 +145,10 @@ class Stoppers: if all(model.gen_state["completed"]): utils.koboldai_vars.generated_tkns = utils.koboldai_vars.genamt del model.gen_state["completed"] + if utils.koboldai_vars.adventure: # Remove added adventure mode stop sequences + for option in extra_options: + if option in utils.koboldai_vars.stop_sequence: + utils.koboldai_vars.stop_sequence.remove(option) return True return False