Update app.py

app.py

@@ -50,22 +50,29 @@ def greedyPredict(input, model, tokenizer):
 def predict_tail(entity, relation):
     global model, tokenizer
     input = entity + "| " + relation
-    out = topkSample(input, model, tokenizer, num_samples=5)
+    out = topkSample(input, model, tokenizer, num_samples=25)
     out_dict = {}
     for k, v in out:
         out_dict[k] = np.exp(v).item()
     return out_dict
 
     
-tokenizer = AutoTokenizer.from_pretrained("apoorvumang/kgt5-wikikg90mv2")
+tokenizer = AutoTokenizer.from_pretrained("apoorvumang/kgt5-base-wikikg90mv2")
 model = AutoModelForSeq2SeqLM.from_pretrained("apoorvumang/kgt5-base-wikikg90mv2")
 
 
-
-ent_input = gr.inputs.Textbox(lines=1, default="World War II")
-rel_input = gr.inputs.Textbox(lines=1, default="followed by")
+ent_input = gr.inputs.Textbox(lines=1, default="Apoorv Umang Saxena")
+rel_input = gr.inputs.Textbox(lines=1, default="country")
 output = gr.outputs.Label()
 
-
-iface = gr.Interface(fn=predict_tail, inputs=[ent_input, rel_input], outputs=output)
+examples = [
+            ['Adrian Kochsiek', 'gender'],
+            ['Apoorv Umang Saxena', 'family name'],
+            ['World War II', 'followed by'],
+            ['Apoorv Umang Saxena', 'country']
+]
+iface = gr.Interface(fn=predict_tail, 
+                     inputs=[ent_input, rel_input], 
+                     outputs=output,
+                     examples=examples,)
 iface.launch()