import ast import copy import random import re import time import traceback import numpy as np import torch import transformers from transformers import LogitsProcessorList import modules.shared as shared from modules.callbacks import ( Iteratorize, Stream, _StopEverythingStoppingCriteria ) from modules.extensions import apply_extensions from modules.html_generator import generate_4chan_html, generate_basic_html from modules.logging_colors import logger from modules.models import clear_torch_cache, local_rank def generate_reply(*args, **kwargs): shared.generation_lock.acquire() try: for result in _generate_reply(*args, **kwargs): yield result finally: shared.generation_lock.release() def get_max_prompt_length(state): return state['truncation_length'] - state['max_new_tokens'] def encode(prompt, add_special_tokens=True, add_bos_token=True, truncation_length=None): if shared.model.__class__.__name__ in ['LlamaCppModel', 'RWKVModel']: input_ids = shared.tokenizer.encode(str(prompt)) input_ids = np.array(input_ids).reshape(1, len(input_ids)) return input_ids else: input_ids = shared.tokenizer.encode(str(prompt), return_tensors='pt', add_special_tokens=add_special_tokens) # This is a hack for making replies more creative. if not add_bos_token and input_ids[0][0] == shared.tokenizer.bos_token_id: input_ids = input_ids[:, 1:] # Handling truncation if truncation_length is not None: input_ids = input_ids[:, -truncation_length:] if shared.model.__class__.__name__ in ['LlamaCppModel', 'RWKVModel', 'ExllamaModel'] or shared.args.cpu: return input_ids elif shared.args.deepspeed: return input_ids.to(device=local_rank) elif torch.backends.mps.is_available(): device = torch.device('mps') return input_ids.to(device) else: return input_ids.cuda() def get_encoded_length(prompt): length_after_extensions = apply_extensions('tokenized_length', prompt) if length_after_extensions is not None: return length_after_extensions return len(encode(prompt)[0]) def decode(output_ids, skip_special_tokens=True): return shared.tokenizer.decode(output_ids, skip_special_tokens) # Removes empty replies from gpt4chan outputs def fix_gpt4chan(s): for i in range(10): s = re.sub("--- [0-9]*\n>>[0-9]*\n---", "---", s) s = re.sub("--- [0-9]*\n *\n---", "---", s) s = re.sub("--- [0-9]*\n\n\n---", "---", s) return s # Fix the LaTeX equations in galactica def fix_galactica(s): s = s.replace(r'\[', r'$') s = s.replace(r'\]', r'$') s = s.replace(r'\(', r'$') s = s.replace(r'\)', r'$') s = s.replace(r'$$', r'$') s = re.sub(r'\n', r'\n\n', s) s = re.sub(r"\n{3,}", "\n\n", s) return s def get_reply_from_output_ids(output_ids, input_ids, original_question, state, is_chat=False): if shared.is_seq2seq: reply = decode(output_ids, state['skip_special_tokens']) else: new_tokens = len(output_ids) - len(input_ids[0]) reply = decode(output_ids[-new_tokens:], state['skip_special_tokens']) # Prevent LlamaTokenizer from skipping a space if type(shared.tokenizer) in [transformers.LlamaTokenizer, transformers.LlamaTokenizerFast] and len(output_ids) > 0: if shared.tokenizer.convert_ids_to_tokens(int(output_ids[-new_tokens])).startswith('▁'): reply = ' ' + reply return reply def formatted_outputs(reply, model_name): if any(s in model_name for s in ['gpt-4chan', 'gpt4chan']): reply = fix_gpt4chan(reply) return reply, generate_4chan_html(reply) else: return reply, generate_basic_html(reply) def set_manual_seed(seed): seed = int(seed) if seed == -1: seed = random.randint(1, 2**31) torch.manual_seed(seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed) return seed def stop_everything_event(): shared.stop_everything = True def generate_reply_wrapper(question, state, stopping_strings=None): reply = question if not shared.is_seq2seq else '' yield formatted_outputs(reply, shared.model_name) for reply in generate_reply(question, state, stopping_strings, is_chat=False): if not shared.is_seq2seq: reply = question + reply yield formatted_outputs(reply, shared.model_name) def apply_stopping_strings(reply, all_stop_strings): stop_found = False for string in all_stop_strings: idx = reply.find(string) if idx != -1: reply = reply[:idx] stop_found = True break if not stop_found: # If something like "\nYo" is generated just before "\nYou:" # is completed, trim it for string in all_stop_strings: for j in range(len(string) - 1, 0, -1): if reply[-j:] == string[:j]: reply = reply[:-j] break else: continue break return reply, stop_found def _generate_reply(question, state, stopping_strings=None, is_chat=False): generate_func = apply_extensions('custom_generate_reply') if generate_func is None: if shared.model_name == 'None' or shared.model is None: logger.error("No model is loaded! Select one in the Model tab.") yield '' return if shared.model.__class__.__name__ in ['LlamaCppModel', 'RWKVModel', 'ExllamaModel']: generate_func = generate_reply_custom else: generate_func = generate_reply_HF # Preparing the input original_question = question if not is_chat: state = apply_extensions('state', state) question = apply_extensions('input', question, state) # Finding the stopping strings all_stop_strings = [] for st in (stopping_strings, ast.literal_eval(f"[{state['custom_stopping_strings']}]")): if type(st) is list and len(st) > 0: all_stop_strings += st if shared.args.verbose: print(f'\n\n{question}\n--------------------\n') shared.stop_everything = False clear_torch_cache() seed = set_manual_seed(state['seed']) last_update = -1 reply = '' is_stream = state['stream'] if len(all_stop_strings) > 0 and not state['stream']: state = copy.deepcopy(state) state['stream'] = True for reply in generate_func(question, original_question, seed, state, stopping_strings, is_chat=is_chat): reply, stop_found = apply_stopping_strings(reply, all_stop_strings) if is_stream: cur_time = time.time() if cur_time - last_update > 0.041666666666666664: # Limit streaming to 24 fps last_update = cur_time yield reply if stop_found: break if not is_chat: reply = apply_extensions('output', reply, state) yield reply def generate_reply_HF(question, original_question, seed, state, stopping_strings=None, is_chat=False): generate_params = {} for k in ['max_new_tokens', 'do_sample', 'temperature', 'top_p', 'typical_p', 'repetition_penalty', 'repetition_penalty_range', 'encoder_repetition_penalty', 'top_k', 'min_length', 'no_repeat_ngram_size', 'num_beams', 'penalty_alpha', 'length_penalty', 'early_stopping', 'tfs', 'top_a', 'mirostat_mode', 'mirostat_tau', 'mirostat_eta']: generate_params[k] = state[k] for k in ['epsilon_cutoff', 'eta_cutoff']: if state[k] > 0: generate_params[k] = state[k] * 1e-4 if state['ban_eos_token']: generate_params['suppress_tokens'] = [shared.tokenizer.eos_token_id] if shared.args.no_cache: generate_params.update({'use_cache': False}) if shared.args.deepspeed: generate_params.update({'synced_gpus': True}) # Encode the input input_ids = encode(question, add_bos_token=state['add_bos_token'], truncation_length=get_max_prompt_length(state)) output = input_ids[0] cuda = not any((shared.args.cpu, shared.args.deepspeed)) if state['auto_max_new_tokens']: generate_params['max_new_tokens'] = state['truncation_length'] - input_ids.shape[-1] # Add the encoded tokens to generate_params question, input_ids, inputs_embeds = apply_extensions('tokenizer', state, question, input_ids, None) original_input_ids = input_ids generate_params.update({'inputs': input_ids}) if inputs_embeds is not None: generate_params.update({'inputs_embeds': inputs_embeds}) # Stopping criteria / eos token eos_token_ids = [shared.tokenizer.eos_token_id] if shared.tokenizer.eos_token_id is not None else [] generate_params['eos_token_id'] = eos_token_ids generate_params['stopping_criteria'] = transformers.StoppingCriteriaList() generate_params['stopping_criteria'].append(_StopEverythingStoppingCriteria()) processor = state.get('logits_processor', LogitsProcessorList([])) # In case folks just pass in a processor by itself. if type(processor) != LogitsProcessorList: processor = LogitsProcessorList([processor]) apply_extensions('logits_processor', processor, input_ids) generate_params['logits_processor'] = processor t0 = time.time() try: if not is_chat and not shared.is_seq2seq: yield '' # Generate the entire reply at once. if not state['stream']: with torch.no_grad(): output = shared.model.generate(**generate_params)[0] if cuda: output = output.cuda() yield get_reply_from_output_ids(output, input_ids, original_question, state, is_chat=is_chat) # Stream the reply 1 token at a time. # This is based on the trick of using 'stopping_criteria' to create an iterator. else: def generate_with_callback(callback=None, *args, **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) with generate_with_streaming(**generate_params) as generator: for output in generator: yield get_reply_from_output_ids(output, input_ids, original_question, state, is_chat=is_chat) if output[-1] in eos_token_ids: break except Exception: traceback.print_exc() finally: t1 = time.time() original_tokens = len(original_input_ids[0]) new_tokens = len(output) - (original_tokens if not shared.is_seq2seq else 0) print(f'Output generated in {(t1-t0):.2f} seconds ({new_tokens/(t1-t0):.2f} tokens/s, {new_tokens} tokens, context {original_tokens}, seed {seed})') return def generate_reply_custom(question, original_question, seed, state, stopping_strings=None, is_chat=False): seed = set_manual_seed(state['seed']) t0 = time.time() reply = '' try: if not is_chat: yield '' if not state['stream']: reply = shared.model.generate(question, state) yield reply else: for reply in shared.model.generate_with_streaming(question, state): yield reply except Exception: traceback.print_exc() finally: t1 = time.time() original_tokens = len(encode(original_question)[0]) new_tokens = len(encode(original_question + reply)[0]) - original_tokens print(f'Output generated in {(t1-t0):.2f} seconds ({new_tokens/(t1-t0):.2f} tokens/s, {new_tokens} tokens, context {original_tokens}, seed {seed})') return