Remove unneeded variables

This commit is contained in:
oobabooga
2023-04-05 22:48:11 -03:00
parent 9a064b78e6
commit f1dd728413

View File

@@ -97,13 +97,8 @@ def extract_message_from_reply(reply, name1, name2, stop_at_newline):
return reply, next_character_found return reply, next_character_found
def chatbot_wrapper(text, generate_params, name1, name2, context, mode, end_of_turn, regenerate=False): def chatbot_wrapper(text, generate_params, name1, name2, context, mode, end_of_turn, regenerate=False):
stop_at_newline = generate_params['stop_at_newline']
max_new_tokens = generate_params['max_new_tokens']
chat_prompt_size = generate_params['chat_prompt_size']
chat_generation_attempts = generate_params['chat_generation_attempts']
just_started = True just_started = True
eos_token = '\n' if stop_at_newline else None eos_token = '\n' if generate_params['stop_at_newline'] else None
name1_original = name1 name1_original = name1
if 'pygmalion' in shared.model_name.lower(): if 'pygmalion' in shared.model_name.lower():
name1 = "You" name1 = "You"
@@ -124,9 +119,9 @@ def chatbot_wrapper(text, generate_params, name1, name2, context, mode, end_of_t
kwargs = {'end_of_turn': end_of_turn, 'is_instruct': mode == 'instruct'} kwargs = {'end_of_turn': end_of_turn, 'is_instruct': mode == 'instruct'}
if custom_generate_chat_prompt is None: if custom_generate_chat_prompt is None:
prompt = generate_chat_prompt(text, max_new_tokens, name1, name2, context, chat_prompt_size, **kwargs) prompt = generate_chat_prompt(text, generate_params['max_new_tokens'], name1, name2, context, generate_params['chat_prompt_size'], **kwargs)
else: else:
prompt = custom_generate_chat_prompt(text, max_new_tokens, name1, name2, context, chat_prompt_size, **kwargs) prompt = custom_generate_chat_prompt(text, generate_params['max_new_tokens'], name1, name2, context, generate_params['chat_prompt_size'], **kwargs)
# Yield *Is typing...* # Yield *Is typing...*
if not regenerate: if not regenerate:
@@ -134,13 +129,13 @@ def chatbot_wrapper(text, generate_params, name1, name2, context, mode, end_of_t
# Generate # Generate
cumulative_reply = '' cumulative_reply = ''
for i in range(chat_generation_attempts): for i in range(generate_params['chat_generation_attempts']):
reply = None reply = None
for reply in generate_reply(f"{prompt}{' ' if len(cumulative_reply) > 0 else ''}{cumulative_reply}", generate_params, eos_token=eos_token, stopping_strings=[f"\n{name1}:", f"\n{name2}:"]): for reply in generate_reply(f"{prompt}{' ' if len(cumulative_reply) > 0 else ''}{cumulative_reply}", generate_params, eos_token=eos_token, stopping_strings=[f"\n{name1}:", f"\n{name2}:"]):
reply = cumulative_reply + reply reply = cumulative_reply + reply
# Extracting the reply # Extracting the reply
reply, next_character_found = extract_message_from_reply(reply, name1, name2, stop_at_newline) reply, next_character_found = extract_message_from_reply(reply, name1, name2, generate_params['stop_at_newline'])
visible_reply = re.sub("(<USER>|<user>|{{user}})", name1_original, reply) visible_reply = re.sub("(<USER>|<user>|{{user}})", name1_original, reply)
visible_reply = apply_extensions(visible_reply, "output") visible_reply = apply_extensions(visible_reply, "output")
@@ -166,27 +161,22 @@ def chatbot_wrapper(text, generate_params, name1, name2, context, mode, end_of_t
yield shared.history['visible'] yield shared.history['visible']
def impersonate_wrapper(text, generate_params, name1, name2, context, mode, end_of_turn): def impersonate_wrapper(text, generate_params, name1, name2, context, mode, end_of_turn):
stop_at_newline = generate_params['stop_at_newline'] eos_token = '\n' if generate_params['stop_at_newline'] else None
max_new_tokens = generate_params['max_new_tokens']
chat_prompt_size = generate_params['chat_prompt_size']
chat_generation_attempts = generate_params['chat_generation_attempts']
eos_token = '\n' if stop_at_newline else None
if 'pygmalion' in shared.model_name.lower(): if 'pygmalion' in shared.model_name.lower():
name1 = "You" name1 = "You"
prompt = generate_chat_prompt(text, max_new_tokens, name1, name2, context, chat_prompt_size, impersonate=True, end_of_turn=end_of_turn) prompt = generate_chat_prompt(text, generate_params['max_new_tokens'], name1, name2, context, generate_params['chat_prompt_size'], impersonate=True, end_of_turn=end_of_turn)
# Yield *Is typing...* # Yield *Is typing...*
yield shared.processing_message yield shared.processing_message
cumulative_reply = '' cumulative_reply = ''
for i in range(chat_generation_attempts): for i in range(generate_params['chat_generation_attempts']):
reply = None reply = None
for reply in generate_reply(f"{prompt}{' ' if len(cumulative_reply) > 0 else ''}{cumulative_reply}", generate_params, eos_token=eos_token, stopping_strings=[f"\n{name1}:", f"\n{name2}:"]): for reply in generate_reply(f"{prompt}{' ' if len(cumulative_reply) > 0 else ''}{cumulative_reply}", generate_params, eos_token=eos_token, stopping_strings=[f"\n{name1}:", f"\n{name2}:"]):
reply = cumulative_reply + reply reply = cumulative_reply + reply
reply, next_character_found = extract_message_from_reply(reply, name1, name2, stop_at_newline) reply, next_character_found = extract_message_from_reply(reply, name1, name2, generate_params['stop_at_newline'])
yield reply yield reply
if next_character_found: if next_character_found:
break break