Spaces:
Runtime error
Runtime error
File size: 2,265 Bytes
fcdf906 d5ad380 fcdf906 cfda2d8 fcdf906 ddcc7db fcdf906 ddcc7db fcdf906 496d05b fcdf906 f81d7bd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 |
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)
|