Implemented temperature scaling and changed output sampling to uniform.

app.py

@@ -71,7 +71,7 @@ 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, :] / 1.2
+        logits = outputs[0, masked_index, :] / temp
         probs = logits.softmax(dim=-1)
         sampling = False
         if sampling:
@@ -114,17 +114,18 @@ PIPELINE_REGISTRY.register_pipeline(
 )
 scrambler = pipeline("temp-scale", model="anferico/bert-for-patents")
 
-def unmask(text, temp):
+def unmask(text, temp, sampling='uniform'):
     # text = add_mask(text)
     split_text = text.split()
-    res = scrambler(text)
+    res = scrambler(text, temp=temp, top_k=3)
     mask_pos = [i for i, t in enumerate(split_text) if 'MASK' in t][0]
     out = {item["token_str"]: item["score"] for item in res}
     score_to_str = {out[k]:k for k in out.keys()}
-    print(score_to_str)
-    print(out)
     score_list = list(score_to_str.keys())
-    idx = np.argmax(np.random.multinomial(1, score_list, 1))
+    if sampling == 'multi':
+        idx = np.argmax(np.random.multinomial(1, score_list, 1))
+    else:
+        idx = np.random.randint(0, len(score_list))
     score = score_list[idx]
     new_token = score_to_str[score]
     split_text[mask_pos] = new_token
@@ -132,7 +133,7 @@ def unmask(text, temp):
 
 textbox = gr.Textbox(label="Type language here", lines=5)
 textbox2 = gr.Textbox(placeholder="Type here...", lines=4)
-temp_slider = gr.Slider(1.0, 1.5, value=1.0, label='Creativity')
+temp_slider = gr.Slider(1.0, 2.0, value=1.0, label='Temperature Scale')
 
 demo = gr.Interface(
     fn=unmask,