Format

app.py

@@ -1,4 +1,5 @@
 import os
+
 os.environ["TOKENIZERS_PARALLELISM"] = "true"
 
 import frontmatter
@@ -12,51 +13,72 @@ from modeling_nova import NovaTokenizer, NovaForCausalLM
 
 print("Downloading model")
 
-tokenizer = AutoTokenizer.from_pretrained('lt-asset/nova-6.7b-bcr', trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained(
+    "lt-asset/nova-6.7b-bcr", trust_remote_code=True
+)
 tokenizer.pad_token = tokenizer.eos_token
 tokenizer.pad_token_id = tokenizer.eos_token_id
 nova_tokenizer = NovaTokenizer(tokenizer)
 
-model = NovaForCausalLM.from_pretrained('lt-asset/nova-6.7b-bcr', torch_dtype=torch.bfloat16, device_map="auto").eval()
+model = NovaForCausalLM.from_pretrained(
+    "lt-asset/nova-6.7b-bcr", torch_dtype=torch.bfloat16, device_map="auto"
+).eval()
 
 examples = json.load(open("humaneval_decompile_nova_6.7b.json", "r"))
 
+
 @spaces.GPU
 def predict(type, normalized_asm):
 
-    prompt_before = f'# This is the assembly code with {type} optimization:\n<func0>:'
+    prompt_before = f"# This is the assembly code with {type} optimization:\n<func0>:"
     asm = normalized_asm.strip()
-    assert asm.startswith('<func0>:')
-    asm = asm[len('<func0>:'): ]
-    prompt_after = '\nWhat is the source code?\n'
-    
+    assert asm.startswith("<func0>:")
+    asm = asm[len("<func0>:") :]
+    prompt_after = "\nWhat is the source code?\n"
+
     inputs = prompt_before + asm + prompt_after
     print("Inputs:", inputs)
 
     # 0 for non-assembly code characters and 1 for assembly characters, required by nova tokenizer
-    char_types = '0' * len(prompt_before) + '1' * len(asm) + '0' * len(prompt_after)
-    
-    tokenizer_output = nova_tokenizer.encode(inputs, '', char_types)
-    input_ids = torch.LongTensor(tokenizer_output['input_ids'].tolist()).unsqueeze(0)
+    char_types = "0" * len(prompt_before) + "1" * len(asm) + "0" * len(prompt_after)
+
+    tokenizer_output = nova_tokenizer.encode(inputs, "", char_types)
+    input_ids = torch.LongTensor(tokenizer_output["input_ids"].tolist()).unsqueeze(0)
     print("Input IDs:", input_ids.shape)
-    nova_attention_mask = torch.LongTensor(tokenizer_output['nova_attention_mask']).unsqueeze(0)
+    nova_attention_mask = torch.LongTensor(
+        tokenizer_output["nova_attention_mask"]
+    ).unsqueeze(0)
 
     output = model.generate(
-        inputs=input_ids.cuda(), max_new_tokens=512, temperature=0.2, top_p=0.95,
-        num_return_sequences=1, do_sample=True, nova_attention_mask=nova_attention_mask.cuda(),
-        no_mask_idx=torch.LongTensor([tokenizer_output['no_mask_idx']]).cuda(), 
-        pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id
+        inputs=input_ids.cuda(),
+        max_new_tokens=512,
+        temperature=0.2,
+        top_p=0.95,
+        num_return_sequences=1,
+        do_sample=True,
+        nova_attention_mask=nova_attention_mask.cuda(),
+        no_mask_idx=torch.LongTensor([tokenizer_output["no_mask_idx"]]).cuda(),
+        pad_token_id=tokenizer.pad_token_id,
+        eos_token_id=tokenizer.eos_token_id,
     )
     print("Output 1:", output)
 
-    output = tokenizer.decode(output[0][input_ids.size(1): ], skip_special_tokens=True, clean_up_tokenization_spaces=True)
+    output = tokenizer.decode(
+        output[0][input_ids.size(1) :],
+        skip_special_tokens=True,
+        clean_up_tokenization_spaces=True,
+    )
     print("Output 2:", output)
 
     return output
 
+
 demo = gr.Interface(
     fn=predict,
-    inputs=[gr.Text(label="Optimization Type", value="O0"), gr.Text(label="Normalized Assembly Code")],
+    inputs=[
+        gr.Text(label="Optimization Type", value="O0"),
+        gr.Text(label="Normalized Assembly Code"),
+    ],
     outputs=gr.Text(label="Raw Nova Output"),
     description=frontmatter.load("README.md").content,
     examples=[[ex["type"], ex["normalized_asm"]] for ex in examples],