import gradio as gr
import spaces
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_path = 'LLM4Binary/llm4decompile-6.7b-v2' # V2 Model
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=torch.bfloat16).cuda()

@spaces.GPU
def predict(input_asm):

    before = f"# This is the assembly code:\n"#prompt
    after = "\n# What is the source code?\n"#prompt
    input_prompt = before+input_asm.strip()+after
    
    inputs = tokenizer(input_prompt, return_tensors="pt").to(model.device)
    with torch.no_grad():
        outputs = model.generate(**inputs, max_new_tokens=2048)### max length to 4096, max new tokens should be below the range
    c_func_decompile = tokenizer.decode(outputs[0][len(inputs[0]):-1])
    return c_func_decompile

demo = gr.Interface(fn=predict,
                    examples=["void ioabs_tcp_pre_select(int *param_1,int *param_2,long param_3) { *param_1 = *param_2; *param_2 = *param_2 + 1; *(int *)((long)*param_1 * 8 + param_3 + 4) = param_1[4]; *(uint *)(param_3 + (long)*param_1 * 8) = *(uint *)(param_3 + (long)*param_1 * 8) | 1; if (((**(int **)(param_1 + 2) + *(int *)(*(long *)(param_1 + 2) + 4)) - *(int *)(*(long *)(param_1 + 2) + 8)) % *(int *)(*(long *)(param_1 + 2) + 4) != 0) { *(uint *)(param_3 + (long)*param_1 * 8) = *(uint *)(param_3 + (long)*param_1 * 8) | 4; } return; }"],
                    inputs="text", outputs="text")
demo.queue()
demo.launch()