Views
No views yet
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer
3from tqdm import tqdm
4
5# Replace 'hf_token' with your Hugging Face token
6hf_token = "your_hf_token_here"
7
8model_name = "ahmedheakl/asm2asm-deepseek1.3b-risc"
9
10instruction = """<|begin▁of▁sentence|>You are a helpful coding assistant assistant on converting from x86 to RISCv64 assembly.
11### Instruction:
12Convert this x86 assembly into RISCv64
13```asm
14{asm_x86}
15"```"
16### Response:
17```asm
18{asm_risc}
19"""
20
21model = AutoModelForCausalLM.from_pretrained(
22 model_name,
23 token=hf_token,
24 device_map="auto",
25 torch_dtype=torch.bfloat16,
26)
27
28model.config.use_cache = True
29
30tokenizer = AutoTokenizer.from_pretrained(
31 model_name,
32 trust_remote_code=True,
33 token=hf_token,
34)
35
36def inference(asm_x86: str) -> str:
37 prompt = instruction.format(asm_x86=asm_x86, asm_risc="")
38 inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
39 generated_ids = model.generate(
40 **inputs,
41 use_cache=True,
42 num_return_sequences=1,
43 max_new_tokens=8000,
44 do_sample=False,
45 num_beams=8,
46 # temperature=0.7,
47 eos_token_id=tokenizer.eos_token_id,
48 pad_token_id=tokenizer.pad_token_id,
49 )
50 outputs = tokenizer.batch_decode(generated_ids)[0]
51 torch.cuda.empty_cache()
52 torch.cuda.synchronize()
53 return outputs.split("```asm\n")[-1].split(f"```{tokenizer.eos_token}")[0]
54
55
56x86 = "DWORD PTR -248[rbp] movsx rdx"
57converted_risc = inference(x86)
58print(converted_risc)| Model | Average Edit Distance (↓) | Exact Match (↑) | Test Accuracy (↑) |
|---|---|---|---|
| GPT4o | 1296 | 0% | 8.18% |
| DeepSeekCoder2-16B | 1633 | 0% | 7.36% |
| Yi-Coder-9B | 1653 | 0% | 6.33% |
| Yi-Coder-1.5B | 275 | 16.98% | 49.69% |
| DeepSeekCoder-1.3B | 107 | 45.91% | 77.23% |
| DeepSeekCoder-1.3B-xTokenizer-int4 | 119 | 46.54% | 72.96% |
| DeepSeekCoder-1.3B-xTokenizer-int8 | 96 | 49.69% | 75.47% |
| DeepSeekCoder-1.3B-xTokenizer | 165 | 50.32% | 79.25% |
| Model | Average Edit Distance (↓) | Exact Match (↑) | Test Accuracy (↑) |
|---|---|---|---|
| GPT4o | 1293 | 0% | 7.55% |
| DeepSeekCoder2-16B | 1483 | 0% | 6.29% |
| DeepSeekCoder-1.3B-xTokenizer-int4 | 112 | 14.47% | 68.55% |
| DeepSeekCoder-1.3B-xTokenizer-int8 | 31 | 69.81% | 88.05% |
| DeepSeekCoder-1.3B-xTokenizer | 27 | 69.81% | 88.68% |
@article{heakl2024cisc,
title={From CISC to RISC: language-model guided assembly transpilation},
author={Heakl, Ahmed and Abi, Chaimaa and Hossam, Rania and Mahmoud, Abdulrahman},
journal={arXiv preprint arXiv:2411.16341},
year={2024}
}