add search type radio

app.py

@@ -33,7 +33,7 @@ model.load_state_dict(torch.load(hf_hub_download(repo_id="nickgardner/chatbot",
                                       filename="alpaca_train_400_epoch.pt"), map_location=device))
 model.eval()
 
-def respond(input):
+def respond(search_type, input):
     model.eval()
     src = torch.tensor(text_pipeline(input), dtype=torch.int64).unsqueeze(0).to(device)
     src_mask = ((src != pad_token) & (src != unknown_token)).unsqueeze(-2).to(device)
@@ -46,18 +46,23 @@ def respond(input):
         trg_mask = torch.autograd.Variable(torch.from_numpy(trg_mask) == 0).to(device)
 
         out = model.out(model.decoder(outputs[:i].unsqueeze(0), e_outputs, src_mask, trg_mask))
-        out = torch.nn.functional.softmax(out, dim=-1)[:, -1].squeeze().detach().numpy()
-        print(out.shape)
-        print(np.sum(out))
-        ix = np.random.choice(np.arange(len(out)), 1, p=out)
-        # val, ix = out[:, -1].data.topk(1)
-
-        # outputs[i] = ix[0][0]
-        outputs[i] = ix[0]
-        # if ix[0][0] == vocab_token_dict['<eos>']:
-        if ix[0] == vocab_token_dict['<eos>']:
-            break
+        if search_type == "Greedy":
+            out = torch.nn.functional.softmax(out, dim=-1)
+            val, ix = out[:, -1].data.topk(1)
+
+            outputs[i] = ix[0][0]
+            if ix[0][0] == vocab_token_dict['<eos>']:
+                break
+        else:
+            out = torch.nn.functional.softmax(out, dim=-1)[:, -1].squeeze().detach().numpy()
+            ix = np.random.choice(np.arange(len(out)), 1, p=out)
+
+            outputs[i] = ix[0]
+            if ix[0] == vocab_token_dict['<eos>']:
+                break
     return ' '.join([indices_to_tokens[ix] for ix in outputs[1:i]])
 
-iface = gr.Interface(fn=respond, inputs="text", outputs="text")
+iface = gr.Interface(fn=respond,
+                     inputs=[gr.Radio(["Greedy", "Probabilistic"], label="Search Type"), "text"],
+                     outputs="text")
 iface.launch()