Update app.py

app.py

@@ -57,7 +57,8 @@ def update_fields(url_or_id):
         return '', '', ''
     print (arxiv_id)
     title, authors, abstract = fetch_arxiv_data(arxiv_id)
-    return title, authors, abstract
+    output = predict(title, authors, abstract)
+    return title, authors, abstract, output
 
 def normalize_spaces(text):
     return re.sub(r'\s+', ' ', text).strip()
@@ -84,7 +85,7 @@ def model_inference(title, authors, abstract):
     score = probs[1].item()
     return score
 
-def predict(title, authors, abstract, _):
+def predict(title, authors, abstract):
     # Your model prediction logic here
     score = model_inference(title, authors, abstract)
     
@@ -96,7 +97,7 @@ def predict(title, authors, abstract, _):
     
     result = f"Your score: {score:.2f}.\nFor papers with score>={threshold:.2f}, {precision * 100:.2f}% are selected by AK.\nFor papers selected by AK, {recall * 100:.2f}% have score>={threshold:.2f}"
     
-    return score, result
+    return result
 
 example_title = "WildChat: 1M ChatGPT Interaction Logs in the Wild"
 example_authors = "Wenting Zhao, Xiang Ren, Jack Hessel, Claire Cardie, Yejin Choi, Yuntian Deng"
@@ -107,18 +108,21 @@ with gr.Blocks() as demo:
     author_box = gr.Textbox(label="Authors (separated by comma)", placeholder="Enter authors (separated by comma)", value=example_authors)
     abstract_box = gr.TextArea(label="Abstract", placeholder="Enter abstract", value=example_abstract)
     arxiv_box = gr.Textbox(label="[Optional] Autofill using arXiv URL/ID", placeholder="[Optional] Autofill using arXiv URL/ID")
+    output_box = gr.Textbox(label="Predicted Selection Probability")
+    autofill_btn = gr.Button("Autofill using arXiv")
+    autofill_btn.click(update_fields, inputs=[arxiv_box], outputs=[title_box, author_box, abstract_box, output_box], concurrency_limit=1)
     #arxiv_box.input(update_fields, inputs=[arxiv_box], outputs=[title_box, author_box, abstract_box])
     
     iface = gr.Interface(
         fn=predict,
-        inputs=[title_box, author_box, abstract_box, arxiv_box],
-        outputs=[gr.Textbox(label="Predicted Score"), gr.Textbox(label="Predicted Selection Probability")],
+        inputs=[title_box, author_box, abstract_box],
+        outputs=[output_box],
+        submit_btn=gr.Button("Predict", variant="primary")
         title="Paper Selection Prediction",
         description="Predict if @_akhaliq will select your paper into Hugging Face papers. Enter the title, authors, and abstract of your paper, or enter an arXiv URL/ID.",
         live=False,
         concurrency_limit=1
     )
-    autofill_btn = gr.Button("Autofill using arXiv")
-    autofill_btn.click(update_fields, inputs=[arxiv_box], outputs=[title_box, author_box, abstract_box], concurrency_limit=1)
+    
 
     demo.queue(max_size=20).launch()