pipeline

app.py

@@ -5,10 +5,12 @@ import re
 import subprocess
 import tempfile
 
+import transformers
+
 import numpy
 print(numpy.__version__)
 
-model = gr.load("ejschwartz/oo-method-test-model-bylibrary", src="models")
+pipe = transformers.pipeline(model="ejschwartz/oo-method-test-model-bylibrary", return_all_scores=True)
 
 # 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',
@@ -101,10 +103,13 @@ with gr.Blocks() as demo:
     def function_change_fn(selected_fun, fun_data):
 
         disassembly_str = fun_data[int(selected_fun, 16)].decode("utf-8")
-        model_results = model(disassembly_str)
+        model_results = pipe(disassembly_str)
         disassembly_str = str(model_results)
 
-        return {disassembly: gr.Textbox.update(value=disassembly_str)}
+        return {disassembly: gr.Textbox.update(value=disassembly_str),
+                #class: gr.Label.update()
+                #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])