let's go field decoding

app.py

@@ -15,10 +15,10 @@ huggingface_hub.login(token=hf_key)
 
 tokenizer = AutoTokenizer.from_pretrained("bigcode/starcoderbase-3b")
 vardecoder_model = AutoModelForCausalLM.from_pretrained(
-    "ejschwartz/resym-vardecoder", torch_dtype=torch.bfloat16#, device_map={"": 0}
+    "ejschwartz/resym-vardecoder", torch_dtype=torch.bfloat16
 ).to("cuda")
 fielddecoder_model = AutoModelForCausalLM.from_pretrained(
-    "ejschwartz/resym-fielddecoder", torch_dtype=torch.bfloat16#, device_map={"": 0}
+    "ejschwartz/resym-fielddecoder", torch_dtype=torch.bfloat16
 ).to("cuda")
 
 gradio_client = Client("https://ejschwartz-resym-field-helper.hf.space/")
@@ -42,10 +42,12 @@ def field_prompt(code):
     print(f"fields: {fields}")
 
     prompt = f"```\n{code}\n```\nWhat are the variable name and type for the following memory accesses:{', '.join(fields)}?\n"
+    if len(fields) > 0:
+        prompt += f"{fields[0]}:"
 
     print(f"field prompt: {prompt}")
 
-    return prompt, field_helper_result
+    return prompt, fields, field_helper_result
 
 @spaces.GPU
 def infer(code):
@@ -65,18 +67,18 @@ def infer(code):
 
     varstring = ", ".join([f"`{v}`" for v in vars])
 
-    var_name = vars[0]
+    first_var = vars[0]
 
     # ejs: Yeah, this var_name thing is really bizarre. But look at https://github.com/lt-asset/resym/blob/main/training_src/fielddecoder_inf.py
-    var_prompt = f"What are the original name and data types of variables {varstring}?\n```\n{code}\n```{var_name}"
+    var_prompt = f"What are the original name and data types of variables {varstring}?\n```\n{code}\n```{first_var}"
 
     print(f"Prompt:\n{var_prompt}")
 
-    input_ids = tokenizer.encode(var_prompt, return_tensors="pt").cuda()[
+    var_input_ids = tokenizer.encode(var_prompt, return_tensors="pt").cuda()[
         :, : 8192 - 1024
     ]
     var_output = vardecoder_model.generate(
-        input_ids=input_ids,
+        input_ids=var_input_ids,
         max_new_tokens=1024,
         num_beams=4,
         num_return_sequences=1,
@@ -86,32 +88,36 @@ def infer(code):
         eos_token_id=0,
     )[0]
     var_output = tokenizer.decode(
-        var_output[input_ids.size(1) :],
+        var_output[var_input_ids.size(1) :],
         skip_special_tokens=True,
         clean_up_tokenization_spaces=True,
     )
 
-    field_prompt_result, field_helper_result = field_prompt(code)
-
-    # field_output = fielddecoder_model.generate(
-    #     input_ids=input_ids,
-    #     max_new_tokens=1024,
-    #     num_beams=4,
-    #     num_return_sequences=1,
-    #     do_sample=False,
-    #     early_stopping=False,
-    #     pad_token_id=0,
-    #     eos_token_id=0,
-    # )[0]
-    # field_output = tokenizer.decode(
-    #     field_output[input_ids.size(1) :],
-    #     skip_special_tokens=True,
-    #     clean_up_tokenization_spaces=True,
-    # )
-
-    var_output = var_name + ":" + var_output
-    # field_output = var_name + ":" + field_output
-    return var_output, varstring
+    field_prompt_result, fields, field_helper_result = field_prompt(code)
+    field_input_ids = tokenizer.encode(field_prompt_result, return_tensors="pt").cuda()[
+        :, : 8192 - 1024
+    ]
+
+    field_output = fielddecoder_model.generate(
+        input_ids=field_input_ids,
+        max_new_tokens=1024,
+        num_beams=4,
+        num_return_sequences=1,
+        do_sample=False,
+        early_stopping=False,
+        pad_token_id=0,
+        eos_token_id=0,
+    )[0]
+    field_output = tokenizer.decode(
+        field_output[var_input_ids.size(1) :],
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=True,
+    )
+
+    var_output = first_var + ":" + var_output
+    if len(fields) > 0:
+        field_output = fields[0] + ":" + field_output
+    return var_output, field_output, varstring
 
 
 demo = gr.Interface(
@@ -121,7 +127,7 @@ demo = gr.Interface(
     ],
     outputs=[
         gr.Text(label="Var Decoder Output"),
-        # gr.Text(label="Field Decoder Output"),
+        gr.Text(label="Field Decoder Output"),
         gr.Text(label="Generated Variable List"),
     ],
     description=frontmatter.load("README.md").content,