Update app.py

app.py

@@ -195,7 +195,10 @@ def get_logits_orig_params_dict(inputs, cond_inputs, alter_label, ex_model, lear
     return params_dict
 
 def edit_process(edit_input, alter_label):
-    _, cond_inputs, edit_inputs, origin_top3 = solve(edit_input, alter_label, edit_task=True)
+    try:
+        _, cond_inputs, edit_inputs, origin_top3 = solve(edit_input, alter_label, edit_task=True)
+    except KeyError:
+        return "The entity or relationship you entered is not in the vocabulary. Please check it carefully.", ""
 
     ### edit output
     fmodel = make_functional(edit_ex_model).eval()
@@ -219,7 +222,10 @@ def edit_process(edit_input, alter_label):
     return "\n".join(origin_top3), "\n".join(edit_top3)
 
 def add_process(edit_input, alter_label):
-    _, cond_inputs, add_inputs, origin_top3 = solve(edit_input, alter_label, edit_task=False)
+    try:
+        _, cond_inputs, add_inputs, origin_top3 = solve(edit_input, alter_label, edit_task=False)
+    except:
+        return "The entity or relationship you entered is not in the vocabulary. Please check it carefully.", ""
 
     ### add output
     fmodel = make_functional(add_ex_model).eval()
@@ -252,7 +258,7 @@ with gr.Blocks() as demo:
         with gr.TabItem("E-FB15k237"):
             with gr.Row():
                 with gr.Column():
-                    edit_input = gr.Textbox(label="Input", lines=1, placeholder="Mask triple input")
+                    edit_input = gr.Textbox(label="Input", lines=1, placeholder=" Please enter in the format of: [MASK]|rel|tail or head|rel|[MASK].")
 
                     alter_label = gr.Textbox(label="Alter Entity", lines=1, placeholder="Entity Name")    
                     edit_button = gr.Button("Edit")
@@ -277,7 +283,7 @@ with gr.Blocks() as demo:
         with gr.TabItem("A-FB15k237"):
             with gr.Row():
                 with gr.Column():
-                    add_input = gr.Textbox(label="Input", lines=1, placeholder="New triple input")
+                    add_input = gr.Textbox(label="Input", lines=1, placeholder="Brand new triple input")
 
                     inductive_entity = gr.Textbox(label="Inductive Entity", lines=1, placeholder="Entity Name")
                     add_button = gr.Button("Add")