use LogitWarpers and add typical_p

app.py

@@ -3,7 +3,15 @@ import torch
 import numpy as np
 import gradio as gr
 from nltk import sent_tokenize
-from transformers import RobertaTokenizer, RobertaForMaskedLM
+
+from transformers import (
+    RobertaTokenizer,
+    RobertaForMaskedLM,
+    LogitsProcessorList,
+    TopKLogitsWarper,
+    TemperatureLogitsWarper,
+)
+from transformers.generation_logits_process import TypicalLogitsWarper
 
 nltk.download('punkt')
 
@@ -17,53 +25,74 @@ if cuda:
 max_len = 20
 top_k = 100
 temperature = 1
+typical_p = 0
 burnin = 250
 max_iter = 500
 
 
 # adapted from https://github.com/nyu-dl/bert-gen
-def generate_step(out,
-                  gen_idx,
-                  temperature=None,
-                  top_k=0,
-                  sample=False,
-                  return_list=True):
+def generate_step(out: object,
+                  gen_idx: int,
+                  top_k: int = top_k,
+                  temperature: float = temperature,
+                  typical_p: float = typical_p,
+                  sample: bool = False) -> list:
     """ Generate a word from from out[gen_idx]
     
     args:
         - out (torch.Tensor): tensor of logits of size batch_size x seq_len x vocab_size
         - gen_idx (int): location for which to generate for
         - top_k (int): if >0, only sample from the top k most probable words
-        - sample (Bool): if True, sample from full distribution. Overridden by top_k 
+        - temperature (float): sampling temperature
+        - typical_p (float): if >0 use typical sampling
+        - sample (bool): if True, sample from full distribution.
+    
+    returns:
+        - list: batch_size tokens
     """
     logits = out.logits[:, gen_idx]
-    if temperature is not None:
-        logits = logits / temperature
+    logit_warpers = []
     if top_k > 0:
-        kth_vals, kth_idx = logits.topk(top_k, dim=-1)
-        dist = torch.distributions.categorical.Categorical(logits=kth_vals)
-        idx = kth_idx.gather(dim=1,
-                             index=dist.sample().unsqueeze(-1)).squeeze(-1)
-    elif sample:
-        dist = torch.distributions.categorical.Categorical(logits=logits)
-        idx = dist.sample()  # removed superfluous squeeze(-1)
+        logit_warpers += [TopKLogitsWarper(top_k)]
+    if temperature:
+        logit_warpers += [TemperatureLogitsWarper(temperature)]
+    if typical_p > 0:
+        if typical_p >= 1:
+            typical_p = 0.999
+        logit_warpers += [TypicalLogitsWarper(typical_p)]
+    logits = LogitsProcessorList(logit_warpers)(None, logits)
+
+    if sample:
+        probs = torch.nn.functional.softmax(logits, dim=-1)
+        next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)
     else:
-        idx = torch.argmax(logits, dim=-1)
-    return idx.tolist() if return_list else idx
+        next_tokens = torch.argmax(logits, dim=-1)
+
+    return next_tokens.tolist()
 
 
 # adapted from https://github.com/nyu-dl/bert-gen
-def parallel_sequential_generation(seed_text,
-                                   seed_end_text,
-                                   max_len=max_len,
-                                   top_k=top_k,
-                                   temperature=temperature,
-                                   max_iter=max_iter,
-                                   burnin=burnin):
-    """ Generate for one random position at a timestep
+def parallel_sequential_generation(seed_text: str,
+                                   seed_end_text: str,
+                                   max_len: int = max_len,
+                                   top_k: int = top_k,
+                                   temperature: float = temperature,
+                                   typical_p: float = typical_p,
+                                   max_iter: int = max_iter,
+                                   burnin: int = burnin) -> str:
+    """ Generate text consistent with preceding and following text
     
-    args:
+    Args:
+        - seed_text (str): preceding text
+        - seed_end_text (str): following text
+        - top_k (int): if >0, only sample from the top k most probable words
+        - temperature (float): sampling temperature
+        - typical_p (float): if >0 use typical sampling
+        - max_iter (int): number of iterations in MCMC
         - burnin: during burn-in period, sample from full distribution; afterwards take argmax
+
+    Returns:
+        - string: generated text to insert between seed_text and seed_end_text
     """
     inp = tokenizer(seed_text + tokenizer.mask_token * max_len + seed_end_text,
                     return_tensors='pt')
@@ -75,12 +104,11 @@ def parallel_sequential_generation(seed_text,
 
     for ii in range(max_iter):
         kk = np.random.randint(0, max_len)
-        out = model(**inp)
-        topk = top_k if (ii >= burnin) else 0
-        idxs = generate_step(out,
+        idxs = generate_step(model(**inp),
                              gen_idx=seed_len + kk,
-                             top_k=topk,
+                             top_k=top_k if (ii >= burnin) else 0,
                              temperature=temperature,
+                             typical_p=typical_p,
                              sample=(ii < burnin))
         inp['input_ids'][0][seed_len + kk] = idxs[0]
 
@@ -90,12 +118,27 @@ def parallel_sequential_generation(seed_text,
     return tokenizer.decode(tokens)
 
 
-def inbertolate(doc,
-                max_len=max_len,
-                top_k=top_k,
-                temperature=temperature,
-                max_iter=max_iter,
-                burnin=burnin):
+def inbertolate(doc: str,
+                max_len: int = max_len,
+                top_k: int = top_k,
+                temperature: float = temperature,
+                typical_p: float = typical_p,
+                max_iter: int = max_iter,
+                burnin: int = burnin):
+    """ Pad out document generating every other sentence
+    
+    Args:
+        - doc (str): document text
+        - max_len (int): number of tokens to insert between sentences
+        - top_k (int): if >0, only sample from the top k most probable words
+        - temperature (float): sampling temperature
+        - typical_p (float): if >0 use typical sampling
+        - max_iter (int): number of iterations in MCMC
+        - burnin: during burn-in period, sample from full distribution; afterwards take argmax
+
+    Returns:
+        - string: generated text to insert between seed_text and seed_end_text
+    """
     new_doc = ''
     paras = doc.split('\n')
 
@@ -108,13 +151,15 @@ def inbertolate(doc,
 
         for sentence in range(len(para) - 1):
             new_doc += para[sentence] + ' '
-            new_doc += parallel_sequential_generation(para[sentence],
-                                                      para[sentence + 1],
-                                                      max_len=max_len,
-                                                      top_k=top_k,
-                                                      temperature=temperature,
-                                                      burnin=burnin,
-                                                      max_iter=max_iter) + ' '
+            new_doc += parallel_sequential_generation(
+                para[sentence],
+                para[sentence + 1],
+                max_len=max_len,
+                top_k=top_k,
+                temperature=float(temperature),
+                typical_p=typical_p,
+                burnin=burnin,
+                max_iter=max_iter) + ' '
 
         new_doc += '\n'
     return new_doc
@@ -130,7 +175,7 @@ if __name__ == '__main__':
         gr.Interface(
             fn=inbertolate,
             inputs=[
-                gr.Textbox(label="Text", lines=7),
+                gr.Textbox(label="Text", lines=10),
                 gr.Slider(label="Maximum length to insert between sentences",
                           minimum=1,
                           maximum=40,
@@ -141,6 +186,10 @@ if __name__ == '__main__':
                           minimum=0,
                           maximum=2,
                           value=temperature),
+                gr.Slider(label="Typical p",
+                          minimum=0,
+                          maximum=1,
+                          value=typical_p),
                 gr.Slider(label="Maximum iterations",
                           minimum=0,
                           maximum=1000,
@@ -150,5 +199,5 @@ if __name__ == '__main__':
                           maximum=500,
                           value=burnin),
             ],
-            outputs=gr.Textbox(label="Expanded text", lines=24))
+            outputs=gr.Textbox(label="Expanded text", lines=30))
     block.launch(server_name='0.0.0.0')