update

main.py

@@ -64,13 +64,13 @@ def compile(compiler, flags, source):
         return None, compile_output, disassembly
 
 
-def disassemble_bytes(byte_data):
+def disassemble_bytes(byte_data, architecture):
     with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as temp_bin_file:
         temp_bin_file.write(byte_data)
         temp_bin_file_name = temp_bin_file.name
 
     disassembly = subprocess.run(
-        ["objdump", "-D", "-b", "binary", "-m", "i386", temp_bin_file_name],
+        ["objdump", "-D", "-b", "binary", "-m", architecture, temp_bin_file_name],
         capture_output=True,
         text=True
     ).stdout
@@ -78,10 +78,10 @@ def disassemble_bytes(byte_data):
     return disassembly
 
 
-def predict(target_bytes, source, compiler, flags):
+def predict(target_bytes, source, compiler, flags, architecture):
     target_bytes = bytes.fromhex(target_bytes)
     compiled_bytes, compile_output, compiled_disassembly = compile(compiler, flags, source)
-    target_disassembly = disassemble_bytes(target_bytes)
+    target_disassembly = disassemble_bytes(target_bytes, architecture)
 
     if compiled_bytes is not None:
         return (
@@ -120,6 +120,7 @@ def run():
             ),
             gr.Textbox(label="Compiler", value="g++"),
             gr.Textbox(label="Compiler Flags", value="-O2"),
+            gr.Textbox(label="Architecture", value="i386"),
         ],
         outputs=[
             gr.Textbox(label="Compiled bytes"),