app.py

import frontmatter
import gradio as gr
import json
import spaces
import torch

from transformers import AutoTokenizer
from modeling_nova import NovaTokenizer, NovaForCausalLM

print("Downloading model")

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()

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>:'
    asm = normalized_asm.strip()
    assert asm.startswith('<func0>:')
    asm = asm[len('<func0>:'): ]
    prompt_after = '\nWhat is the source code?\n'
    
    inputs = prompt_before + asm + prompt_after
    # 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)
    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=20, 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
    )

    output = tokenizer.decode(output[input_ids.size(1): ], skip_special_tokens=True, clean_up_tokenization_spaces=True)

    return output

demo = gr.Interface(
    fn=predict,
    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],
)
demo.launch(show_error=True)