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

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

description = f"""
# LLM4Decompile 9B V2

This is a space for testing the [LLM4Decompile 9B V2 model](https://huggingface.co/LLM4Binary/llm4decompile-9b-v2).  It expects to be given a decompiled function output by Ghidra.  I simply copy and paste from the Ghidra GUI, but this is not the method recommended by the official model page, which could influence performance.  I am not affiliated with the authors.
"""

@spaces.GPU(duration=120)
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,
                    description=description,
                    examples=["""undefined4 func0(float param_1,long param_2,int param_3)
{
  int local_28;
  int local_24;
  
  local_24 = 0;
  do {
    local_28 = local_24;
    if (param_3 <= local_24) {
      return 0;
    }
    while (local_28 = local_28 + 1, local_28 < param_3) {
      if ((double)((ulong)(double)(*(float *)(param_2 + (long)local_24 * 4) -
                                  *(float *)(param_2 + (long)local_28 * 4)) &
                  SUB168(_DAT_00402010,0)) < (double)param_1) {
        return 1;
      }
    }
    local_24 = local_24 + 1;
  } while( true );
}""","""
undefined8 __cdecl FUN_004025f0(byte *param_1,int param_2)

{
  do {
    *param_1 = ~*param_1;
    *param_1 = *param_1 - 0x37;
    *param_1 = *param_1 + 1;
    param_1 = param_1 + 1;
    param_2 = param_2 + -1;
  } while (param_2 != 0);
  return 0x3632000137;
}

""","""
void __thiscall FUN_00404f30(void *this,char param_1,uint param_2)

{
  char *pcVar1;
  
  if ((param_1 != '\0') && (0xf < *(uint *)((int)this + 0x18))) {
    pcVar1 = *(char **)((int)this + 4);
    if (param_2 != 0) {
      std::_Traits_helper::copy_s<>((char *)((int)this + 4),0x10,pcVar1,param_2);
    }
    FUN_00405320(pcVar1);
  }
  *(undefined4 *)((int)this + 0x18) = 0xf;
  std::basic_string<>::_Eos((basic_string<> *)this,param_2);
  return;
}

""","""
char __thiscall
FUN_008a28f0(void *param_1_00,FMOD_VECTOR *param_1,undefined4 ****param_2,undefined4 param_3,
            undefined4 param_4,float param_5)

{
  char cVar1;
  undefined4 uVar2;
  undefined4 ****ppppuVar3;
  char *unaff_EBX;
  FMOD_VECTOR *pFVar4;
  char *pcVar5;
  undefined4 uVar6;
  FMOD_VECTOR *local_a8;
  undefined4 ***local_a4;
  undefined4 local_a0;
  uint uStack_90;
  basic_string<> abStack_8c [4];
  basic_string<> local_88 [28];
  basic_string<> abStack_6c [4];
  undefined4 ***apppuStack_68 [5];
  uint local_54 [8];
  basic_string<> abStack_34 [28];
  void *local_18;
  int iStack_14;
  undefined4 uStack_10;
  void *pvStack_c;
  undefined *puStack_8;
  undefined4 uStack_4;
  
  uStack_4 = 0xffffffff;
  puStack_8 = &LAB_009a39d4;
  pvStack_c = ExceptionList;
  if (param_1 == (FMOD_VECTOR *)0x0) {
    if (param_1_00 != (void *)0x0) {
      ExceptionList = &pvStack_c;
      std::basic_string<>::~basic_string<>((basic_string<> *)((int)param_1_00 + 4));
                    /* WARNING: Subroutine does not return */
      operator_delete(param_1_00);
    }
    cVar1 = '\0';
  }
  else {
    if (param_1_00 == (void *)0x0) {
      ExceptionList = &pvStack_c;
      if (0.1 < param_5) {
        ExceptionList = &pvStack_c;
        FUN_00736f20();
        FUN_008a0dc0();
        FMOD::Event::setCallback
                  ((_func_FMOD_RESULT_FMOD_EVENT_ptr_FMOD_EVENT_CALLBACKTYPE_void_ptr_void_ptr_void_ ptr
                    *)param_1,&LAB_008a0fb0);
      }
    }
    else {
      ExceptionList = &pvStack_c;
      uVar2 = FUN_00736f20();
      uVar6 = 1;
      pcVar5 = 
      "Manager::PlaySoundEvent(); attempting to register FMOD callback for 2D one-shot Event.";
      *(undefined4 *)((int)param_1_00 + 0x24) = uVar2;
      FMOD::Event::setCallback
                ((_func_FMOD_RESULT_FMOD_EVENT_ptr_FMOD_EVENT_CALLBACKTYPE_void_ptr_void_ptr_void_pt r
                  *)param_1,&LAB_008a0fe0);
      FUN_008a0a60(param_1_00,pcVar5,uVar6);
    }
    uVar2 = 0;
    local_a8 = param_1;
    local_a4 = param_2;
    local_a0 = param_3;
    FMOD::Event::set3DAttributes(param_1,(FMOD_VECTOR *)&local_a8,(FMOD_VECTOR *)0x0);
    cVar1 = FUN_008a0a60("Manager::PlaySoundEvent(); attempting to set 3D attributes for a 3D one-sh ot Event."
                         ,1,uVar2);
    if (cVar1 == '\0') {
      uVar2 = 0;
      FMOD::Event::getInfo((int *)param_1,(char **)0x0,(FMOD_EVENT_INFO *)&stack0xffffff4c);
      std::basic_string<>::basic_string<>(local_88,unaff_EBX);
      uStack_10 = 4;
      uVar2 = std::operator+<>(abStack_34,
                               "Manager::PlaySoundEvent: unable to set3DAttributes for sound event ( "
                               ,local_88,uVar2);
      uStack_10._0_1_ = 5;
      std::operator+<>(abStack_6c,uVar2,&DAT_00aa30ec);
      uStack_10._0_1_ = 7;
      std::basic_string<>::~basic_string<>(abStack_34);
      if (local_54[0] < 0x10) {
        apppuStack_68[0] = apppuStack_68;
      }
      FUN_008a0b10(apppuStack_68[0]);
      uStack_10 = CONCAT31(uStack_10._1_3_,4);
      std::basic_string<>::~basic_string<>(abStack_6c);
      uStack_10 = 0xffffffff;
      std::basic_string<>::~basic_string<>(local_88);
      cVar1 = '\0';
    }
    else {
      pFVar4 = param_1;
      FMOD::Event::start();
      cVar1 = FUN_008a0a60("Manager::PlaySoundEvent(); attempting to play 3D one-shot Event.",1,
                           pFVar4);
      if (cVar1 == '\0') {
        uVar2 = 0;
        FMOD::Event::getInfo((int *)param_1,(char **)0x0,(FMOD_EVENT_INFO *)&stack0xffffff4c);
        std::basic_string<>::basic_string<>((basic_string<> *)local_54,unaff_EBX);
        iStack_14 = 0;
        uVar2 = std::operator+<>(abStack_8c,"Manager::PlaySoundEvent: Failed to start sound event ("
                                 ,local_54,uVar2);
        iStack_14._0_1_ = 1;
        std::operator+<>(&local_a8,uVar2,&DAT_00aa30ec);
        iStack_14._0_1_ = 3;
        std::basic_string<>::~basic_string<>(abStack_8c);
        ppppuVar3 = (undefined4 ****)local_a4;
        if (uStack_90 < 0x10) {
          ppppuVar3 = &local_a4;
        }
        FUN_008a0b10(ppppuVar3);
        iStack_14 = (uint)iStack_14._1_3_ << 8;
        std::basic_string<>::~basic_string<>((basic_string<> *)&local_a8);
        iStack_14 = 0xffffffff;
        std::basic_string<>::~basic_string<>((basic_string<> *)local_54);
        cVar1 = '\0';
      }
      else if (0.1 < (float)param_2) {
        FUN_00736f20();
        FUN_008a0dc0();
      }
    }
  }
  ExceptionList = local_18;
  return cVar1;
}

"""],
                    inputs="text", outputs="text")
demo.queue()
demo.launch()