Fix merge conflict in text_generation
- Need to update `shared.still_streaming = False` before the final `yield formatted_outputs`, shifted the position of some yields.
This commit is contained in:
Xan
commit
b3e10e47c0
10 changed files with 298 additions and 155 deletions
|
|
@ -5,13 +5,13 @@ import time
|
|||
import numpy as np
|
||||
import torch
|
||||
import transformers
|
||||
from tqdm import tqdm
|
||||
|
||||
import modules.shared as shared
|
||||
from modules.callbacks import (Iteratorize, Stream,
|
||||
_SentinelTokenStoppingCriteria)
|
||||
from modules.extensions import apply_extensions
|
||||
from modules.html_generator import generate_4chan_html, generate_basic_html
|
||||
from modules.models import local_rank
|
||||
from modules.stopping_criteria import _SentinelTokenStoppingCriteria
|
||||
|
||||
|
||||
def get_max_prompt_length(tokens):
|
||||
|
|
@ -92,19 +92,22 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi
|
|||
# These models are not part of Hugging Face, so we handle them
|
||||
# separately and terminate the function call earlier
|
||||
if shared.is_RWKV:
|
||||
if shared.args.no_stream:
|
||||
reply = shared.model.generate(context=question, token_count=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k)
|
||||
yield formatted_outputs(reply, shared.model_name)
|
||||
else:
|
||||
yield formatted_outputs(question, shared.model_name)
|
||||
# RWKV has proper streaming, which is very nice.
|
||||
# No need to generate 8 tokens at a time.
|
||||
for reply in shared.model.generate_with_streaming(context=question, token_count=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k):
|
||||
try:
|
||||
if shared.args.no_stream:
|
||||
reply = shared.model.generate(context=question, token_count=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k)
|
||||
yield formatted_outputs(reply, shared.model_name)
|
||||
|
||||
t1 = time.time()
|
||||
print(f"Output generated in {(t1-t0):.2f} seconds.")
|
||||
return
|
||||
else:
|
||||
yield formatted_outputs(question, shared.model_name)
|
||||
# RWKV has proper streaming, which is very nice.
|
||||
# No need to generate 8 tokens at a time.
|
||||
for reply in shared.model.generate_with_streaming(context=question, token_count=max_new_tokens, temperature=temperature, top_p=top_p, top_k=top_k):
|
||||
yield formatted_outputs(reply, shared.model_name)
|
||||
finally:
|
||||
t1 = time.time()
|
||||
output = encode(reply)[0]
|
||||
input_ids = encode(question)
|
||||
print(f"Output generated in {(t1-t0):.2f} seconds ({(len(output)-len(input_ids[0]))/(t1-t0):.2f} tokens/s, {len(output)-len(input_ids[0])} tokens)")
|
||||
return
|
||||
|
||||
original_question = question
|
||||
if not (shared.args.chat or shared.args.cai_chat):
|
||||
|
|
@ -113,23 +116,19 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi
|
|||
print(f"\n\n{question}\n--------------------\n")
|
||||
|
||||
input_ids = encode(question, max_new_tokens)
|
||||
original_input_ids = input_ids
|
||||
output = input_ids[0]
|
||||
cuda = "" if (shared.args.cpu or shared.args.deepspeed or shared.args.flexgen) else ".cuda()"
|
||||
n = shared.tokenizer.eos_token_id if eos_token is None else int(encode(eos_token)[0][-1])
|
||||
stopping_criteria_list = transformers.StoppingCriteriaList()
|
||||
if stopping_string is not None:
|
||||
# The stopping_criteria code below was copied from
|
||||
# https://github.com/PygmalionAI/gradio-ui/blob/master/src/model.py
|
||||
# Copied from https://github.com/PygmalionAI/gradio-ui/blob/master/src/model.py
|
||||
t = encode(stopping_string, 0, add_special_tokens=False)
|
||||
stopping_criteria_list = transformers.StoppingCriteriaList([
|
||||
_SentinelTokenStoppingCriteria(
|
||||
sentinel_token_ids=t,
|
||||
starting_idx=len(input_ids[0])
|
||||
)
|
||||
])
|
||||
else:
|
||||
stopping_criteria_list = None
|
||||
stopping_criteria_list.append(_SentinelTokenStoppingCriteria(sentinel_token_ids=t, starting_idx=len(input_ids[0])))
|
||||
|
||||
if not shared.args.flexgen:
|
||||
generate_params = [
|
||||
f"max_new_tokens=max_new_tokens",
|
||||
f"eos_token_id={n}",
|
||||
f"stopping_criteria=stopping_criteria_list",
|
||||
f"do_sample={do_sample}",
|
||||
|
|
@ -147,45 +146,23 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi
|
|||
]
|
||||
else:
|
||||
generate_params = [
|
||||
f"max_new_tokens={max_new_tokens if shared.args.no_stream else 8}",
|
||||
f"do_sample={do_sample}",
|
||||
f"temperature={temperature}",
|
||||
f"stop={n}",
|
||||
]
|
||||
if shared.args.deepspeed:
|
||||
generate_params.append("synced_gpus=True")
|
||||
if shared.args.no_stream:
|
||||
generate_params.append("max_new_tokens=max_new_tokens")
|
||||
else:
|
||||
generate_params.append("max_new_tokens=8")
|
||||
if shared.soft_prompt:
|
||||
inputs_embeds, filler_input_ids = generate_softprompt_input_tensors(input_ids)
|
||||
generate_params.insert(0, "inputs_embeds=inputs_embeds")
|
||||
generate_params.insert(0, "filler_input_ids")
|
||||
generate_params.insert(0, "inputs=filler_input_ids")
|
||||
else:
|
||||
generate_params.insert(0, "input_ids")
|
||||
|
||||
# Generate the entire reply at once
|
||||
if shared.args.no_stream:
|
||||
with torch.no_grad():
|
||||
output = eval(f"shared.model.generate({', '.join(generate_params)}){cuda}")[0]
|
||||
if shared.soft_prompt:
|
||||
output = torch.cat((input_ids[0], output[filler_input_ids.shape[1]:]))
|
||||
|
||||
reply = decode(output)
|
||||
if not (shared.args.chat or shared.args.cai_chat):
|
||||
reply = original_question + apply_extensions(reply[len(question):], "output")
|
||||
|
||||
t1 = time.time()
|
||||
print(f"Output generated in {(t1-t0):.2f} seconds ({(len(output)-len(input_ids[0]))/(t1-t0)/8:.2f} it/s, {len(output)-len(input_ids[0])} tokens)")
|
||||
yield formatted_outputs(reply, shared.model_name)
|
||||
|
||||
# Generate the reply 8 tokens at a time
|
||||
else:
|
||||
yield formatted_outputs(original_question, shared.model_name)
|
||||
shared.still_streaming = True
|
||||
for i in tqdm(range(max_new_tokens//8+1)):
|
||||
clear_torch_cache()
|
||||
generate_params.insert(0, "inputs=input_ids")
|
||||
|
||||
try:
|
||||
# Generate the entire reply at once.
|
||||
if shared.args.no_stream:
|
||||
with torch.no_grad():
|
||||
output = eval(f"shared.model.generate({', '.join(generate_params)}){cuda}")[0]
|
||||
if shared.soft_prompt:
|
||||
|
|
@ -194,22 +171,66 @@ def generate_reply(question, max_new_tokens, do_sample, temperature, top_p, typi
|
|||
reply = decode(output)
|
||||
if not (shared.args.chat or shared.args.cai_chat):
|
||||
reply = original_question + apply_extensions(reply[len(question):], "output")
|
||||
|
||||
if not shared.args.flexgen:
|
||||
if output[-1] == n:
|
||||
break
|
||||
input_ids = torch.reshape(output, (1, output.shape[0]))
|
||||
else:
|
||||
if np.count_nonzero(input_ids[0] == n) < np.count_nonzero(output == n):
|
||||
break
|
||||
input_ids = np.reshape(output, (1, output.shape[0]))
|
||||
|
||||
#Mid-stream yield, ran if no breaks
|
||||
|
||||
yield formatted_outputs(reply, shared.model_name)
|
||||
|
||||
if shared.soft_prompt:
|
||||
inputs_embeds, filler_input_ids = generate_softprompt_input_tensors(input_ids)
|
||||
|
||||
#Stream finished from max tokens or break. Do final yield.
|
||||
shared.still_streaming = False
|
||||
yield formatted_outputs(reply, shared.model_name)
|
||||
# Stream the reply 1 token at a time.
|
||||
# This is based on the trick of using 'stopping_criteria' to create an iterator.
|
||||
elif not shared.args.flexgen:
|
||||
|
||||
def generate_with_callback(callback=None, **kwargs):
|
||||
kwargs['stopping_criteria'].append(Stream(callback_func=callback))
|
||||
clear_torch_cache()
|
||||
with torch.no_grad():
|
||||
shared.model.generate(**kwargs)
|
||||
|
||||
def generate_with_streaming(**kwargs):
|
||||
return Iteratorize(generate_with_callback, kwargs, callback=None)
|
||||
|
||||
shared.still_streaming = True
|
||||
yield formatted_outputs(original_question, shared.model_name)
|
||||
with eval(f"generate_with_streaming({', '.join(generate_params)})") as generator:
|
||||
for output in generator:
|
||||
if shared.soft_prompt:
|
||||
output = torch.cat((input_ids[0], output[filler_input_ids.shape[1]:]))
|
||||
reply = decode(output)
|
||||
|
||||
if not (shared.args.chat or shared.args.cai_chat):
|
||||
reply = original_question + apply_extensions(reply[len(question):], "output")
|
||||
|
||||
if output[-1] == n:
|
||||
break
|
||||
yield formatted_outputs(reply, shared.model_name)
|
||||
|
||||
shared.still_streaming = False
|
||||
yield formatted_outputs(reply, shared.model_name)
|
||||
|
||||
# Stream the output naively for FlexGen since it doesn't support 'stopping_criteria'
|
||||
else:
|
||||
shared.still_streaming = True
|
||||
for i in range(max_new_tokens//8+1):
|
||||
clear_torch_cache()
|
||||
with torch.no_grad():
|
||||
output = eval(f"shared.model.generate({', '.join(generate_params)})")[0]
|
||||
if shared.soft_prompt:
|
||||
output = torch.cat((input_ids[0], output[filler_input_ids.shape[1]:]))
|
||||
reply = decode(output)
|
||||
|
||||
if not (shared.args.chat or shared.args.cai_chat):
|
||||
reply = original_question + apply_extensions(reply[len(question):], "output")
|
||||
|
||||
if np.count_nonzero(input_ids[0] == n) < np.count_nonzero(output == n):
|
||||
break
|
||||
yield formatted_outputs(reply, shared.model_name)
|
||||
|
||||
input_ids = np.reshape(output, (1, output.shape[0]))
|
||||
if shared.soft_prompt:
|
||||
inputs_embeds, filler_input_ids = generate_softprompt_input_tensors(input_ids)
|
||||
|
||||
shared.still_streaming = False
|
||||
yield formatted_outputs(reply, shared.model_name)
|
||||
|
||||
finally:
|
||||
t1 = time.time()
|
||||
print(f"Output generated in {(t1-t0):.2f} seconds ({(len(output)-len(original_input_ids[0]))/(t1-t0):.2f} tokens/s, {len(output)-len(original_input_ids[0])} tokens)")
|
||||
return
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue