Use Luke's clean_output function

main.py

@@ -26,6 +26,56 @@ def canonicalize_output(decompiled: str):
     assert False, "Unable to parse code"
   return df.canonical_code + DECOMPILED_ORIG_SEP
 
+# From Luke to clean up the results
+from prepare import (
+    parser, 
+    FileTypeMapping, 
+    PreprocessedFunction, 
+    has_error, 
+    TypeNotDefinedError, 
+    TypeNotFoundError, 
+    UnsupportedFeatureError,
+    get_all_user_defined_types
+)
+from evaluator import add_placeholders_for_nonvariable_types
+from tree_sitter import Node
+
+class InvalidPredictionError(Exception):
+    pass
+
+def clean_output(raw_prediction: str) -> str:
+    ### Parse and sort nodes
+    root = parser.parse(bytes(raw_prediction, "utf8")).root_node
+    if root.type == "ERROR":
+        raise InvalidPredictionError("Prediction is irrecoverably malformed.")
+    assert root.type == "translation_unit"
+    fn_nodes: list[Node] = []
+    udt_nodes: list[Node] = []
+    other_nodes: list[Node] = []
+    for node in root.children:
+        if node.type == "function_definition":
+            fn_nodes.append(node)
+        elif node.type in {"struct_specifier", "union_specifier", "enum_specifier"}:
+            udt_nodes.append(node)
+        elif node.type == ";":
+            pass
+        else:
+            other_nodes.append(node)
+
+    # Process all UDTs
+    types = FileTypeMapping()
+    for type_node in udt_nodes:
+        if not has_error(type_node):
+            try:
+                types.parse_type(type_node)
+            except:
+                pass # Ignore the misgenerated type. If it's necessary, we'll get another exception later.
+
+    if len(fn_nodes) == 0:
+        raise InvalidPredictionError("No recognizable function definitions.") 
+    fn_node = fn_nodes[0]
+    assert fn_node.text is not None # To make mypy happy
+
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline #, BitsAndBytesConfig
 from peft import PeftModel
 
@@ -88,7 +138,7 @@ def predict(inp):
     except Exception as e:
         return f"Something went wrong during inference: {str(e)}"
     print("done!")
-    return results[0]['generated_text']
+    return clean_output(results[0]['generated_text'])
 
 def run():