Added multinomial prediction sampling.

app.py

@@ -20,7 +20,7 @@ def add_mask(text, size=1):
 
 
 class TempScalePipe(FillMaskPipeline):
-    def postprocess(self, model_outputs, top_k=5, target_ids=None):
+    def postprocess(self, model_outputs, top_k=3, target_ids=None):
         # Cap top_k if there are targets
         if target_ids is not None and target_ids.shape[0] < top_k:
             top_k = target_ids.shape[0]
@@ -30,8 +30,10 @@ class TempScalePipe(FillMaskPipeline):
         masked_index = torch.nonzero(input_ids == self.tokenizer.mask_token_id, as_tuple=False).squeeze(-1)
         # Fill mask pipeline supports only one ${mask_token} per sample
 
-        logits = outputs[0, masked_index, :] / 1e3
+        logits = outputs[0, masked_index, :] / 1e1
         probs = logits.softmax(dim=-1)
+        indices = torch.multinomial(probs, num_samples=3)
+        probs = probs[indices]
         if target_ids is not None:
             probs = probs[..., target_ids]
 
@@ -98,7 +100,8 @@ from transformers import pipeline, Pipeline
 # textbox = gr.Textbox(label="Type language here", lines=5)
 #
 demo = gr.Interface(
-    fn=unmask,
+    unmask,
+    [gr.Slider(minimum=0, maximum=15, value=8, step=1, label="Guidance scale")],
     inputs=textbox,
     outputs="label",
     examples=[example],