def instructions(): if settings.getboolean("console-bell"): bell = "on" else: bell = "off" if settings.getboolean("action-d20"): d20 = "on" else: d20 = "off" print( '\033[' + colors["instructions"] + 'm' + 'AID2: Clover Edition Instructions: \n Enter actions starting with a verb ex. "go to the tavern" or "attack the orc."\n To speak enter say "(thing you want to say)" or just "(thing you want to say)"' ) print('The following commands can be entered for any action:') print( ' "/revert" Reverts the last action allowing you to pick a different action.' ) print(' "/quit" Quits the game and saves') print( ' "/menu" Starts a new game and saves your current one') print(' "/retry" Retries the last action') print(' "/restart" Restarts the current story') print( ' "/print" Prints a transcript of your adventure (without extra newline formatting)' ) print(' "/help" Prints these instructions again') print( ' "/set SETTING VALUE" Sets the specified setting to the specified value.:' ) print( ' temp Higher values make the AI more random. Default: 0.4 | Current:', settings.getfloat("temp")) print( ' rep-pen Controls how repetitive the AI is allowed to be. Default: 1.2 | Current:', settings.getfloat("rep-pen")) print( ' text-wrap-width Maximum width of lines printed by computer. Default: 80 | Current:', settings.getint("text-wrap-width")) print( ' console-bell Beep after AI generates text? Default: on | Current:', bell) print( ' top-keks Number of words the AI can randomly choose. Default: 20 | Current:', settings.getint("top-keks")) print(' generate-num Default: 60 | Current:', settings.getint("generate-num")) print(' top-p Default: 0.9 | Current:', settings.getfloat("top-p")) print(' log-level Default: 3 | Current:', settings.getint("log-level")) print( ' action-sugg How many actions to generate, 0 is off. Default: 4 | Current:', settings.getint("action-sugg")) print( ' action-d20 Make actions difficult. Default: on | Current:', d20) print( ' action-temp How random the suggested actions are. Default: 1 | Current:', settings.getfloat("action-temp"), '\033[39m')
def clear_lines(n): """Clear the last line in the terminal.""" if in_colab() or settings.getboolean('colab-mode'): # this wont work in colab etc return screen_code = "\033[1A[\033[2K" # up one line, and clear line for _ in range(n): print(screen_code, end="\r")
def __init__(self, generate_num=60, temperature=0.4, top_k=40, top_p=0.9, dtype=DTYPE, model_path: Union[str, Path] = Path('models', 'pytorch-gpt2-xl-aid2-v5'), repetition_penalty=1, repetition_penalty_range=512, repetition_penalty_slope=3.33): self.generate_num = generate_num self.temp = temperature self.top_k = top_k self.top_p = top_p self.samples = 1 self.dtype = dtype self.repetition_penalty = repetition_penalty self.repetition_penalty_range = repetition_penalty_range self.repetition_penalty_slope = repetition_penalty_slope self.batch_size = 1 self.max_history_tokens = 1024 - generate_num self.stop_token = "<|endoftext|>" if isinstance(model_path, str): self.checkpoint_path = model_path logger.warning( f"Using DEBUG MODE! This will load one of the generic (non-finetuned) GPT2 models. " f"Selected: {model_path}") elif isinstance(model_path, Path): self.checkpoint_path = model_path if not self.checkpoint_path.exists(): raise FileNotFoundError( "Could not find {} Make sure to download a pytorch model and put it in the models directory!" .format(str(self.checkpoint_path))) else: raise ValueError( f"model_path must be either str or Path, got {type(model_path)}" ) self.device = torch.device("cuda" if self.dtype == torch.float16 else "cpu") logger.info("Using device={}, checkpoint={}, dtype={}".format( self.device, str(self.checkpoint_path), self.dtype)) # Load tokenizer and model model_class, tokenizer_class = MODEL_CLASSES[ "gpt2-experimental"] if settings.getboolean( "gpt2_experimental") else MODEL_CLASSES["gpt2"] if "gpt-neo" in str(model_path): self.max_history_tokens = 2048 - generate_num model_class = GPTNeoForCausalLM self.tokenizer = tokenizer_class.from_pretrained( str(self.checkpoint_path)) self.model = model_class.from_pretrained(str(self.checkpoint_path)) self.model.to(self.dtype).to(self.device) self.model.eval()
def sample_sequence(model, length, context, temperature=1, top_k=0, top_p=0.9, repetition_penalty=1.0, device="cpu", stop_tokens=None, tokenizer=None): """Actually generate the tokens""" logger.debug('temp: {} top_k: {} top_p: {} rep-pen: {}'.format( temperature, top_k, top_p, repetition_penalty)) max_length = context.shape[1] + length # check to see if greater than 2048? if settings.getboolean('force-cpu'): context = context.long().cpu() else: context = context.long().cuda() out = model.generate( context, do_sample=True, min_length=max_length, max_length=max_length, temperature=temperature, top_k=top_k, top_p=top_p, repetition_penalty=repetition_penalty, repetition_penalty_range=300, repetition_penalty_slope=3.33, use_cache=True, pad_token_id=tokenizer.eos_token_id, ).long() generated = tokenizer.decode(out[0]) return generated
def in_colab(): """Some terminal codes don't work in a colab notebook.""" # from https://github.com/tqdm/tqdm/blob/master/tqdm/autonotebook.py if settings.getboolean("colab-mode"): return True try: from IPython import get_ipython if (not get_ipython()) or ( 'IPKernelApp' not in get_ipython().config): # pragma: no cover raise ImportError("console") if 'VSCODE_PID' in os.environ: # pragma: no cover raise ImportError("vscode") except ImportError: if get_terminal_size()[0] == 0 or 'google.colab' in sys.modules: settings["colab-mode"] = "on" settings["prompt-toolkit"] = "off" return True return False else: settings["colab-mode"] = "on" settings["prompt-toolkit"] = "off" return True
def play(generator): print("\n") with open(Path("interface", "mainTitle.txt"), "r", encoding="utf-8") as file: colPrint(file.read(), colors["title"], wrap=False) with open(Path("interface", "subTitle.txt"), "r", encoding="utf-8") as file: cols = termWidth for line in file: line = re.sub(r'\n', '', line) line = line[:cols] #fills in the graphic using reverse video mode substituted into the areas between |'s colPrint( re.sub(r'\|[ _]*(\||$)', lambda x: '\x1B[7m' + x.group(0) + '\x1B[27m', line), colors['subtitle'], False) print() colPrint( "Go to https://github.com/cloveranon/Clover-Edition/ or email [email protected] for bug reports, help, and feature requests.", colors['subsubtitle']) while True: # May be needed to avoid out of mem gc.collect() torch.cuda.empty_cache() print("\n\n") colPrint( "0: Pick Prompt From File (Default if you type nothing)\n1: Write Custom Prompt", colors['menu']) if getNumberInput(1) == 1: with open(Path("interface", "prompt-instructions.txt"), "r", encoding="utf-8") as file: colPrint(file.read(), colors["instructions"], False) prompt = colInput("Prompt>", colors["main-prompt"], colors["user-text"]) context = colInput("Context>", colors["main-prompt"], colors["user-text"]) filename = colInput( "Name to save prompt as? (Leave blank for no save): ", colors["query"], colors["user-text"], ) filename = re.sub( "-$", "", re.sub("^-", "", re.sub("[^a-zA-Z0-9_-]+", "-", filename))) if filename != "": with open(Path("prompts", filename + ".txt"), "w", encoding="utf-8") as f: f.write(context + "\n" + prompt) else: prompt, context = selectFile() assert (prompt + context) instructions() print() colPrint("Generating story...", colors["loading-message"]) story = newStory(generator, prompt, context) while True: # Generate suggested actions act_alts = settings.getint("action-sugg") if act_alts > 0: # TODO change this to two messages for different colors suggested_actions = [] colPrint("\nSuggested actions:", colors["selection-value"]) action_suggestion_lines = 2 for i in range(act_alts): suggested_action = story.getSuggestion() if len(suggested_action.strip()) > 0: j = len(suggested_actions) suggested_actions.append(suggested_action) suggestion = "{}> {}".format(j, suggested_action) action_suggestion_lines += colPrint( suggestion, colors["selection-value"]) print() bell() action = colInput("> You ", colors["main-prompt"], colors["user-text"]) # Clear suggestions and user input if act_alts > 0: action_suggestion_lines += 2 if not IN_COLAB: clear_lines(action_suggestion_lines) # Show user input again # colPrint("\n> " + action.rstrip(), colors["user-text"], end="") setRegex = re.search("^/set ([^ ]+) ([^ ]+)$", action) if setRegex: if setRegex.group(1) in settings: currentSettingValue = settings[setRegex.group(1)] colPrint( "Current Value of {}: {} Changing to: {}".format( setRegex.group(1), currentSettingValue, setRegex.group(2))) settings[setRegex.group(1)] = setRegex.group(2) colPrint("Save config file?", colors["query"]) colPrint("Saving an invalid option will corrupt file!", colors["error"]) if (colInput( "y/n? >", colors["selection-prompt"], colors["selection-value"], ) == "y"): with open("config.ini", "w", encoding="utf-8") as file: config.write(file) else: colPrint("Invalid Setting", colors["error"]) instructions() elif action == "/menu": break elif action == "/restart": print() colPrint("Restarting story...", colors["loading-message"]) story = newStory(generator, story.prompt, context) continue elif action == "/quit": exit() elif action == "/help": instructions() elif action == "/print": print("\nPRINTING\n") #TODO colorize printed story colPrint(str(story), colors["print-story"]) elif action == '/retry': if len(story.story) == 1: print() colPrint("Restarting story...", colors["loading-message"]) story = newStory(generator, story.prompt, context) continue else: newaction = story.story[-1][0] colPrint(newaction, colors['user-text'], end='') story.story = story.story[:-1] result = "\n" + story.act(newaction)[0] if len(story.story) >= 2: similarity = get_similarity(result, story.story[-2][1][0]) if similarity > 0.9: story.story = story.story[:-1] colPrint( "Woops that action caused the model to start looping. Try a different action to prevent that.", colors["error"], ) continue colPrint(result, colors["ai-text"]) continue elif action == '/revert': if len(story.story) == 1: colPrint("You can't go back any farther. ", colors["error"]) continue story.story = story.story[:-1] colPrint("Last action reverted. ", colors["message"]) if len(story.story) < 2: colPrint(story.prompt, colors["ai-text"]) colPrint(story.story[-1][1][0], colors["ai-text"]) continue elif action == "/alter": story.story[-1][1][0] = alterText(story.story[-1][1][0]) if len(story.story) < 2: colPrint(story.prompt, colors["ai-text"]) else: colPrint("\n" + story.story[-1][0] + "\n", colors["transformed-user-text"]) colPrint("\n" + story.story[-1][1][0] + "\n\n", colors["ai-text"]) elif action == "/prompt": story.prompt = alterText(story.prompt) if len(story.story) < 2: colPrint(story.prompt, colors["ai-text"]) else: colPrint("\n" + story.story[-1][0] + "\n", colors["transformed-user-text"]) colPrint("\n" + story.story[-1][1][0] + "\n\n", colors["ai-text"]) else: if act_alts > 0: # Options to select a suggestion action if action in [ str(i) for i in range(len(suggested_actions)) ]: action = suggested_actions[int(action)] original_action = action action = action.strip() #TODO debug stuff to delete if action != original_action: logger.debug("STRIPPED WHITE SPACE OFF ACTION %r vs %r", original_action, action) # Crop actions to a max length #action = action[:4096] if action != "": # Roll a 20 sided dice to make things interesting d = random.randint(1, 20) logger.debug("roll d20=%s", d) # If it says 'You say "' then it's still dialouge. Normalise it by removing `You say `, we will add again soon action = re.sub("^ ?[Yy]ou say [\"']", '"', action) if any(action.lstrip().startswith(t) for t in ['"', "'"]): if settings.getboolean("action-d20"): action = d20ify_speech(action, d) else: action = "You say " + action logger.info( "%r. %r, %r", action, any(action.lstrip().startswith(t) for t in ['"', "'"]), settings.getboolean("action-d20")) else: action = first_to_second_person(action) if not action.lower().startswith( "you ") and not action.lower().startswith( "i "): action = action[0].lower() + action[1:] # roll a d20 if settings.getboolean("action-d20"): action = d20ify_action(action, d) else: action = "You " + action if action[-1] not in [".", "?", "!"]: action = action + "." action = "\n> " + action + "\n" colPrint( "\n>" + action.lstrip().lstrip("> \n"), colors["transformed-user-text"], ) #TODO check if leading white space makes sense result = "\n" + story.act(action)[0] #TODO: Replace all this nonsense if len(story.story) >= 2: similarity = get_similarity(result, story.story[-2][1][0]) if similarity > 0.9: story.story = story.story[:-1] colPrint( "Woops that action caused the model to start looping. Try a different action to prevent that.", colors["error"], ) continue if player_won(result): colPrint(result + "\n CONGRATS YOU WIN", colors["message"]) break elif player_died(result): colPrint(result, colors["ai-text"]) colPrint("YOU DIED. GAME OVER", colors["error"]) colPrint( "\nOptions:\n0)Start a new game\n1)\"I'm not dead yet!\" (If you didn't actually die)", colors["menu"], ) choice = getNumberInput(1) if choice == 0: break else: colPrint("Sorry about that...where were we?", colors["query"]) colPrint(result, colors["ai-text"])
def bell(): if settings.getboolean("console-bell"): print("\x07", end="")
if (not get_ipython()) or ( 'IPKernelApp' not in get_ipython().config): # pragma: no cover raise ImportError("console") if 'VSCODE_PID' in os.environ: # pragma: no cover raise ImportError("vscode") except ImportError: if get_terminal_size()[0] == 0 or 'google.colab' in sys.modules: return True return False else: return True IN_COLAB = _in_colab() logger.info("Colab detected: {}".format(IN_COLAB)) IN_COLAB = IN_COLAB or settings.getboolean('colab-mode') if IN_COLAB: logger.warning( "Colab mode enabled, disabling line clearing and readline to avoid colab bugs." ) else: try: import readline logger.info( 'readline has been imported. This enables a number of editting features but may cause bugs for colab users.' ) except ModuleNotFoundError: pass termWidth = get_terminal_size()[0] if termWidth < 5:
def sample_sequence(model, length, context, temperature=1, top_k=0, top_p=0.9, repetition_penalty=1.0, repetition_penalty_range=512, repetition_penalty_slope=3.33, device="cpu", stop_tokens=None, tokenizer=None): """Actually generate the tokens""" logger.debug( 'temp: {} top_k: {} top_p: {} rep-pen: {} rep-pen-range: {} rep-pen-slope: {}' .format(temperature, top_k, top_p, repetition_penalty, repetition_penalty_range, repetition_penalty_slope)) context_tokens = context context = torch.tensor(context, dtype=torch.long, device=device) # context = context.repeat(num_samples, 1) generated = context USE_PAST = True next_token = context pasts = None clines = 0 penalty = None if not repetition_penalty_range is None and not repetition_penalty_slope is None and repetition_penalty_range > 0: penalty = (torch.arange(repetition_penalty_range) / (repetition_penalty_range - 1)) * 2. - 1 penalty = (repetition_penalty_slope * penalty) / (1 + torch.abs(penalty) * (repetition_penalty_slope - 1)) penalty = 1 + ((penalty + 1) / 2) * (repetition_penalty - 1) with torch.no_grad(): for j in range(length): # why would we ever not use past? # is generated and next_token always same thing? if not USE_PAST: input_ids_next = generated pasts = None else: input_ids_next = next_token # Note: we could also use 'past' with GPT-2/Transfo-XL/XLNet/CTRL (cached hidden-states) model_kwargs = {"past": pasts, "use_cache": True} model_inputs = model.prepare_inputs_for_generation( generated.unsqueeze(0), **model_kwargs) model_outputs = model(**model_inputs, return_dict=True) logits, pasts = model_outputs.logits, model_outputs.past_key_values logits = logits[0, -1, :].float() # Originally the order was Temperature, Repetition Penalty, then top-k/p if settings.getboolean('top-p-first'): logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p) logits = logits / (temperature if temperature > 0 else 1.0) # repetition penalty from CTRL (https://arxiv.org/abs/1909.05858) plus range limit if repetition_penalty != 1.0: if penalty is not None: penalty_len = min(generated.shape[0], repetition_penalty_range) penalty_context = generated[-repetition_penalty_range:] score = torch.gather(logits, 0, penalty_context) penalty = penalty.type(score.dtype).to(score.device) penalty_window = penalty[-penalty_len:] score = torch.where(score < 0, score * penalty_window, score / penalty_window) logits.scatter_(0, penalty_context, score) else: score = torch.gather(logits, 0, generated) score = torch.where(score < 0, score * repetition_penalty, score / repetition_penalty) logits.scatter_(0, generated, score) if not settings.getboolean('top-p-first'): logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p) if temperature == 0: # greedy sampling: next_token = torch.argmax(logits, dim=-1).unsqueeze(-1) else: next_token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1) generated = torch.cat((generated, next_token), dim=-1) # Decode into plain text o = generated[len(context_tokens):].tolist() generated.text = tokenizer.decode( o, clean_up_tokenization_spaces=False, skip_special_tokens=True) if use_ptoolkit(): clear_lines(clines) generated.text = format_result(generated.text) clines = output(generated.text, "ai-text") if ((stop_tokens is not None) and (j > 4) and (next_token[0] in stop_tokens)): # Why the minimum tokens, j>X. Because sometimes the models starts with whitespace, which will strip away anyway. Having a minimum amount of tokens before we stop usually means we don't just stop because of "\n " or similar logger.debug( "Stopping generation as we found stop tokens. One of `%s`, in '%s'. token generated `%s`", stop_tokens, next_token, j, ) break clear_lines(clines) return generated
import os from pathlib import Path from typing import Union import torch import torch.nn.functional as F import re from gpt2 import GPT2LMHeadModelExperimental from transformers import GPT2Tokenizer, GPT2LMHeadModel, GPTNeoForCausalLM from getconfig import settings, logger from utils import cut_trailing_sentence, output, clear_lines, format_result, use_ptoolkit if not settings.getboolean('force-cpu') and not torch.cuda.is_available(): logger.warning('CUDA is not available, you are limited to CPU only.') DTYPE = torch.float32 if ((not torch.cuda.is_available()) or settings.getboolean('force-cpu')) else torch.float16 logger.info('Cuda Available: {} Force CPU: {} Precision: {}'.format( torch.cuda.is_available(), settings.getboolean('force-cpu'), '32-bit' if DTYPE == torch.float32 else '16-bit')) # warnings.filterwarnings("ignore") MODEL_CLASSES = { "gpt2": (GPT2LMHeadModel, GPT2Tokenizer), "gpt2-experimental": (GPT2LMHeadModelExperimental, GPT2Tokenizer), } def getTokens(tokenizer, l): tokenizer.encode()
def sample_sequence( model, length, context, temperature=1, top_k=0, top_p=0.9, repetition_penalty=1.0, device="cpu", stop_tokens=None, tokenizer=None ): """Actually generate the tokens""" logger.debug( 'temp: {} top_k: {} top_p: {} rep-pen: {}'.format(temperature, top_k, top_p, repetition_penalty)) context_tokens = context context = torch.tensor(context, dtype=torch.long, device=device) # context = context.repeat(num_samples, 1) generated = context USE_PAST = True next_token = context pasts = None clines = 0 with torch.no_grad(): for j in range(length): # why would we ever not use past? # is generated and next_token always same thing? if not USE_PAST: input_ids_next = generated pasts = None else: input_ids_next = next_token # Note: we could also use 'past' with GPT-2/Transfo-XL/XLNet/CTRL (cached hidden-states) logits, pasts = model(input_ids=input_ids_next, past=pasts) logits = logits[-1, :].float() # переписать логику TODO if settings.getboolean('sparse-gen'): probs = entmax_bisect(logits, dim=-1, alpha=settings.sparse-level) next_token = torch.multinomial(probs, num_samples=1) else: # Originally the order was Temperature, Repetition Penalty, then top-k/p if settings.getboolean('top-p-first'): logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p) logits = logits / (temperature if temperature > 0 else 1.0) # repetition penalty from CTRL (https://arxiv.org/abs/1909.05858) for k in set(generated.tolist()): logits[k] /= repetition_penalty if not settings.getboolean('top-p-first'): logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p) if temperature == 0: # greedy sampling: next_token = torch.argmax(logits, dim=-1).unsqueeze(-1) else: next_token = torch.multinomial( F.softmax(logits, dim=-1), num_samples=1 ) generated = torch.cat((generated, next_token), dim=-1) # Decode into plain text o = generated[len(context_tokens):].tolist() generated.text = tokenizer.decode( o, clean_up_tokenization_spaces=False, skip_special_tokens=True ) if use_ptoolkit(): clear_lines(clines) generated.text = format_result(generated.text) clines = output(generated.text, "ai-text") if ( (stop_tokens is not None) and (j > 4) and (next_token[0] in stop_tokens) ): # Why the minimum tokens, j>X. Because sometimes the models starts with whitespace, which will strip away anyway. Having a minimum amount of tokens before we stop usually means we don't just stop because of "\n " or similar logger.debug( "Stopping generation as we found stop tokens. One of `%s`, in '%s'. token generated `%s`", stop_tokens, next_token, j, ) break clear_lines(clines) return generated
import os from pathlib import Path from typing import Union import torch import torch.nn.functional as F import re from gpt2 import GPT2LMHeadModelExperimental from transformers import GPT2Tokenizer, GPT2LMHeadModel from getconfig import settings, logger from utils import cut_trailing_sentence, output, clear_lines, format_result, use_ptoolkit if not settings.getboolean('force-cpu') and not torch.cuda.is_available(): logger.warning('CUDA is not available, you are limited to CPU only.') DTYPE = torch.float32 if ((not torch.cuda.is_available()) or settings.getboolean('force-cpu')) else torch.float16 logger.info('Cuda Available: {} Force CPU: {} Precision: {}'.format(torch.cuda.is_available(), settings.getboolean('force-cpu'), '32-bit' if DTYPE == torch.float32 else '16-bit')) # warnings.filterwarnings("ignore") MODEL_CLASSES = { "gpt2": (GPT2LMHeadModel, GPT2Tokenizer), "gpt2-experimental": (GPT2LMHeadModelExperimental, GPT2Tokenizer), } def getTokens(tokenizer, l): tokenizer.encode()
def use_ptoolkit(): return not settings.getboolean("colab-mode") and settings.getboolean( 'prompt-toolkit')
"""Clear the last line in the terminal.""" if in_colab() or settings.getboolean('colab-mode'): # this wont work in colab etc return screen_code = "\033[1A[\033[2K" # up one line, and clear line for _ in range(n): print(screen_code, end="\r") if in_colab(): logger.warning( "Colab mode enabled, disabling line clearing and readline to avoid colab bugs." ) else: try: if settings.getboolean('prompt-toolkit'): from inline_editor import edit_multiline from prompt_toolkit import prompt as ptprompt from prompt_toolkit import print_formatted_text from prompt_toolkit.styles import Style from prompt_toolkit.formatted_text import to_formatted_text, HTML else: raise ModuleNotFoundError logger.info( 'Python Prompt Toolkit has been imported. This enables a number of editing features but may cause bugs for colab users.' ) except (ImportError, ModuleNotFoundError): try: settings['prompt-toolkit'] = "off" import readline
def sample_sequence(model, length, context, num_samples=1, temperature=1, top_k=0, top_p=0.9, repetition_penalty=1.0, is_xlnet=False, is_xlm_mlm=False, xlm_mask_token=None, xlm_lang=None, device="cpu", stop_tokens=None, tokenizer=None): logger.debug('temp: {} top_k: {} top_p: {} rep-pen: {}'.format( temperature, top_k, top_p, repetition_penalty)) context = torch.tensor(context, dtype=torch.long, device=device) context = context.unsqueeze(0).repeat(num_samples, 1) generated = context USE_PAST = True next_token = context outputs = None with torch.no_grad(): for j in range(length): #why would we ever not use past? #is generated and next_token always same thing? if USE_PAST: past = outputs[1] if outputs is not None else None inputs = {"input_ids": next_token, "past": past} else: inputs = {"input_ids": generated} outputs = model( **inputs ) # Note: we could also use 'past' with GPT-2/Transfo-XL/XLNet/CTRL (cached hidden-states) logits = outputs[0][:, -1, :].float() #Originally the order was Temperature, Repetition Penalty, then top-k/p if settings.getboolean('top-p-first'): logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p) logits = logits / (temperature if temperature > 0 else 1.0) # repetition penalty from CTRL (https://arxiv.org/abs/1909.05858) for i in range(num_samples): for k in set(generated[i].tolist()): logits[i, k] /= repetition_penalty if not settings.getboolean('top-p-first'): logits = top_k_top_p_filtering(logits, top_k=top_k, top_p=top_p) if temperature == 0: # greedy sampling: next_token = torch.argmax(logits, dim=-1).unsqueeze(-1) else: next_token = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1) generated = torch.cat((generated, next_token), dim=1) if ((stop_tokens is not None) and (j > 4) and (next_token[0][0] in stop_tokens)): # Why the minimum tokens, j>X. Because sometimes the models starts with whitespace, which will strip away anyway. Having a minimum amount of tokens before we stop usually means we don't just stop because of "\n " or similar logger.debug( "Stopping generation as we found stop tokens. One of `%s`, in '%s'. token generated `%s`", stop_tokens, next_token, j, ) break return generated
import os from pathlib import Path import itertools import torch import torch.nn.functional as F from transformers import GPT2LMHeadModel, GPT2Tokenizer from getconfig import settings, logger from story.utils import cut_trailing_sentence DTYPE = torch.float32 if ((not torch.cuda.is_available()) or settings.getboolean('force-cpu')) else torch.float16 logger.info('Cuda Available: {} Force CPU: {} DTYPE: {}'.format( torch.cuda.is_available(), settings.getboolean('force-cpu'), DTYPE)) # warnings.filterwarnings("ignore") MODEL_CLASSES = { "gpt2": (GPT2LMHeadModel, GPT2Tokenizer), } def top_k_top_p_filtering(logits, top_k=0, top_p=0.0, filter_value=-float("Inf")): """ Filter a distribution of logits using top-k and/or nucleus (top-p) filtering Args: logits: logits distribution shape (batch size x vocabulary size) top_k > 0: keep only top k tokens with highest probability (top-k filtering). top_p > 0.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
def use_ptoolkit(): return not in_colab() and settings.getboolean('prompt-toolkit')