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()