Spaces:
Runtime error
Runtime error
Stop generation at \nHuman and \n---- in Chat mode (#15)
Browse files- Stop generation at at `\nHuman` and `\n----` in chat mode (76bb25ab1928288a5620db6a484d51ff7966172d)
Co-authored-by: Loubna Ben Allal <loubnabnl@users.noreply.huggingface.co>
app.py
CHANGED
|
@@ -89,6 +89,8 @@ def generate(prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition
|
|
| 89 |
do_sample=True,
|
| 90 |
seed=42,
|
| 91 |
)
|
|
|
|
|
|
|
| 92 |
|
| 93 |
if chat_mode and FIM_INDICATOR in prompt:
|
| 94 |
raise ValueError("Chat mode and FIM are mutually exclusive. Choose one or the other.")
|
|
@@ -114,11 +116,15 @@ def generate(prompt, temperature=0.9, max_new_tokens=256, top_p=0.95, repetition
|
|
| 114 |
else:
|
| 115 |
output = prompt
|
| 116 |
|
|
|
|
| 117 |
for response in stream:
|
| 118 |
if fim_mode and response.token.text =="<|endoftext|>":
|
| 119 |
output += (suffix + "\n" + response.token.text)
|
|
|
|
|
|
|
| 120 |
else:
|
| 121 |
output += response.token.text
|
|
|
|
| 122 |
yield output
|
| 123 |
return output
|
| 124 |
|
|
|
|
| 89 |
do_sample=True,
|
| 90 |
seed=42,
|
| 91 |
)
|
| 92 |
+
if chat_mode:
|
| 93 |
+
generate_kwargs.update({"stop_sequences": ["\nHuman", "\n-----"]})
|
| 94 |
|
| 95 |
if chat_mode and FIM_INDICATOR in prompt:
|
| 96 |
raise ValueError("Chat mode and FIM are mutually exclusive. Choose one or the other.")
|
|
|
|
| 116 |
else:
|
| 117 |
output = prompt
|
| 118 |
|
| 119 |
+
previous_token = ""
|
| 120 |
for response in stream:
|
| 121 |
if fim_mode and response.token.text =="<|endoftext|>":
|
| 122 |
output += (suffix + "\n" + response.token.text)
|
| 123 |
+
elif chat_mode and response.token.text in ["Human", "-----"] and previous_token=="\n":
|
| 124 |
+
return output
|
| 125 |
else:
|
| 126 |
output += response.token.text
|
| 127 |
+
previous_token = response.token.text
|
| 128 |
yield output
|
| 129 |
return output
|
| 130 |
|