More sophisticated attempt at fixing tabs

app.py

@@ -13,6 +13,44 @@ from normalize import normalize
 from transformers import AutoTokenizer
 from modeling_nova import NovaTokenizer, NovaForCausalLM
 
+
+def fix_assembly_tabs(asm_text):
+    """
+    Fix assembly code formatting by ensuring proper tab placement.
+    Expected format: address:TABhex_bytesWHITESPACEinstructionWHITESPACEoperands
+    """
+    lines = asm_text.split("\n")
+    fixed_lines = []
+
+    for line in lines:
+        line = line.rstrip()  # Remove trailing whitespace
+        if not line.strip():  # Skip empty lines
+            fixed_lines.append(line)
+            continue
+
+        # Check if this looks like an assembly instruction line
+        # Pattern: optional_spaces + hex_address + colon + hex_bytes + instruction + operands
+        asm_pattern = r"^(\s*)([0-9a-f]+):\s*([0-9a-f\s]+?)\s+(\w+)(\s+.*)?$"
+        match = re.match(asm_pattern, line, re.IGNORECASE)
+
+        if match:
+            indent, address, hex_bytes, instruction, operands = match.groups()
+            operands = operands or ""
+
+            # Clean up hex bytes (remove extra spaces)
+            hex_bytes = re.sub(r"\s+", " ", hex_bytes.strip())
+
+            # Reconstruct with proper tab formatting
+            # Format: indent + address + ":" + TAB + hex_bytes + TAB + instruction + operands
+            fixed_line = f"{indent}{address}:\t{hex_bytes}\t{instruction}{operands}"
+            fixed_lines.append(fixed_line)
+        else:
+            # Not an assembly instruction line, keep as is
+            fixed_lines.append(line)
+
+    return "\n".join(fixed_lines)
+
+
 print("Downloading model")
 
 tokenizer = AutoTokenizer.from_pretrained(
@@ -33,6 +71,7 @@ examples = json.load(open("humaneval_decompile_nova_6.7b.json", "r"))
 def predict(type, input_asm, _c_source):
 
     if "<func0>:" not in input_asm:
+        # Needs normalizing
 
         # Add a bogus function header if needed.
         first_line = input_asm.split("\n")[0]
@@ -40,8 +79,9 @@ def predict(type, input_asm, _c_source):
             print("Adding synthetic function header")
             input_asm = "<func0>:\n" + input_asm
 
-        # Convert consecutive spaces to tabs
-        input_asm = re.sub(r' {2,}', '\t', input_asm)
+        # Fix tab formatting in assembly code
+        input_asm = fix_assembly_tabs(input_asm)
+
         normalized_asm = normalize(input_asm)
     else:
         normalized_asm = input_asm