class

app.py

@@ -10,7 +10,7 @@ import transformers
 import numpy
 print(numpy.__version__)
 
-pipe = transformers.pipeline(model="ejschwartz/oo-method-test-model-bylibrary", return_all_scores=True)
+pipe = transformers.pipeline(model="ejschwartz/oo-method-test-model-bylibrary", top_k=2)
 
 # gr.load("ejschwartz/oo-method-test-model-bylibrary", src="models", examples=[
 #     '\nL1: 55 ?? push ebp\n 8b ec ?? mov ebp, esp\n 5d ?? pop ebp\n c3 ?? ret\n\n',
@@ -80,7 +80,7 @@ with gr.Blocks() as demo:
         fun_dropdown = gr.Dropdown(label="Select a function", choices=["Woohoo!"], interactive=True)
 
         with gr.Row(visible=True) as result:
-            disassembly = gr.Textbox(label="Disassembly", lines=20, readonly=True)
+            disassembly = gr.Textbox(label="Disassembly", lines=20)
             clazz = gr.Label()
 
     def file_change_fn(file):
@@ -104,15 +104,16 @@ with gr.Blocks() as demo:
 
         disassembly_str = fun_data[int(selected_fun, 16)].decode("utf-8")
         model_results = pipe(disassembly_str)
-        disassembly_str = str(model_results)
+        top_k = {i["label"].split(", ")[0]: i["score"] for i in model_results}
+        #disassembly_str = str(model_results)
 
         return {disassembly: gr.Textbox.update(value=disassembly_str),
-                #class: gr.Label.update()
+                clazz: gr.Label.update(top_k)
                 #clazz: gr.Label.update(value="Method" if model_results[0]["label"] == "method" else "Function")}
         }
 
     file_widget.change(file_change_fn, file_widget, [col, fun_dropdown, all_dis_state])
 
-    fun_dropdown.change(function_change_fn, [fun_dropdown, all_dis_state], [disassembly])
+    fun_dropdown.change(function_change_fn, [fun_dropdown, all_dis_state], [disassembly, clazz])
 
 demo.launch(server_name="0.0.0.0", server_port=7860, share=True)