ejschwartz commited on
Commit
3f47af7
·
1 Parent(s): 894ff3a

Try using their own code

Browse files
Files changed (2) hide show
  1. app.py +10 -0
  2. prep_decompiled.py +0 -0
app.py CHANGED
@@ -7,6 +7,8 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
7
 
8
  import huggingface_hub
9
 
 
 
10
  print("Hello!")
11
 
12
  hf_key = os.environ["HF_TOKEN"]
@@ -19,6 +21,8 @@ vardecoder_model = AutoModelForCausalLM.from_pretrained(
19
  "ejschwartz/resym-vardecoder", torch_dtype=torch.bfloat16, device_map="auto"
20
  )
21
 
 
 
22
  example = """{
23
  "input": "What are the original name and data type of variables `a1`, `a2`, `k`, `j`, `i`?\n```\n_BYTE *__fastcall sub_4022CD(_BYTE *a1, __int64 a2)\n{\n_BYTE *result; // rax\n__int16 v4; // [rsp+1Ch] [rbp-14h]\nunsigned __int16 v5; // [rsp+1Eh] [rbp-12h]\nunsigned __int16 v6; // [rsp+20h] [rbp-10h]\nunsigned __int16 v7; // [rsp+22h] [rbp-Eh]\nunsigned int k; // [rsp+24h] [rbp-Ch]\nunsigned int j; // [rsp+28h] [rbp-8h]\nunsigned int i; // [rsp+2Ch] [rbp-4h]\n\nfor ( i = 0; i <= 2; ++i )\n{\nfor ( j = 0; j <= 0x3F; ++j )\n{\nfor ( k = 0; k <= 3; ++k )\n{\n*(&v4 + k) = *(_WORD *)(a2 + 2 * (k + 4 * j + ((unsigned __int64)i << 8)));\n*(&v4 + k) += (*(&v4 + k) >> 15) & 0xD01;\n*(&v4 + k) = ((((unsigned __int16)*(&v4 + k) << 10) + 1664) / 0xD01u) & 0x3FF;\n}\n*a1 = v4;\na1[1] = (4 * v5) | HIBYTE(v4);\na1[2] = (16 * v6) | (v5 >> 6);\na1[3] = ((_BYTE)v7 << 6) | (v6 >> 4);\nresult = a1 + 4;\na1[4] = v7 >> 2;\na1 += 5;\n}\n}\nreturn result;\n}\n```",
24
  "output": "a1: r, uint8_t*\na2: a, const polyvec*\nk: t, uint16_t\nj: -, -\ni: k, unsigned int",
@@ -38,6 +42,12 @@ example = """{
38
 
39
  @spaces.GPU
40
  def infer(var_name, code):
 
 
 
 
 
 
41
  #line = json.loads(input)
42
  #first_token = line["output"].split(":")[0]
43
  prompt = code + var_name + ":"
 
7
 
8
  import huggingface_hub
9
 
10
+ import prep_decompiled
11
+
12
  print("Hello!")
13
 
14
  hf_key = os.environ["HF_TOKEN"]
 
21
  "ejschwartz/resym-vardecoder", torch_dtype=torch.bfloat16, device_map="auto"
22
  )
23
 
24
+
25
+
26
  example = """{
27
  "input": "What are the original name and data type of variables `a1`, `a2`, `k`, `j`, `i`?\n```\n_BYTE *__fastcall sub_4022CD(_BYTE *a1, __int64 a2)\n{\n_BYTE *result; // rax\n__int16 v4; // [rsp+1Ch] [rbp-14h]\nunsigned __int16 v5; // [rsp+1Eh] [rbp-12h]\nunsigned __int16 v6; // [rsp+20h] [rbp-10h]\nunsigned __int16 v7; // [rsp+22h] [rbp-Eh]\nunsigned int k; // [rsp+24h] [rbp-Ch]\nunsigned int j; // [rsp+28h] [rbp-8h]\nunsigned int i; // [rsp+2Ch] [rbp-4h]\n\nfor ( i = 0; i <= 2; ++i )\n{\nfor ( j = 0; j <= 0x3F; ++j )\n{\nfor ( k = 0; k <= 3; ++k )\n{\n*(&v4 + k) = *(_WORD *)(a2 + 2 * (k + 4 * j + ((unsigned __int64)i << 8)));\n*(&v4 + k) += (*(&v4 + k) >> 15) & 0xD01;\n*(&v4 + k) = ((((unsigned __int16)*(&v4 + k) << 10) + 1664) / 0xD01u) & 0x3FF;\n}\n*a1 = v4;\na1[1] = (4 * v5) | HIBYTE(v4);\na1[2] = (16 * v6) | (v5 >> 6);\na1[3] = ((_BYTE)v7 << 6) | (v6 >> 4);\nresult = a1 + 4;\na1[4] = v7 >> 2;\na1 += 5;\n}\n}\nreturn result;\n}\n```",
28
  "output": "a1: r, uint8_t*\na2: a, const polyvec*\nk: t, uint16_t\nj: -, -\ni: k, unsigned int",
 
42
 
43
  @spaces.GPU
44
  def infer(var_name, code):
45
+
46
+ splitcode = code.splitlines()
47
+ comments = prep_decompiled.extract_comments(splitcode)
48
+ sig = prep_decompiled.parse_signature(splitcode)
49
+ print(f"comments={comments} sig={sig}")
50
+
51
  #line = json.loads(input)
52
  #first_token = line["output"].split(":")[0]
53
  prompt = code + var_name + ":"
prep_decompiled.py ADDED
The diff for this file is too large to render. See raw diff