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-xtokenizer-arm"
9
10instruction = """<|begin▁of▁sentence|>You are a helpful coding assistant specialized in converting from x86 to ARM assembly.
11### Instruction:
12Convert this x86 assembly into ARM
13```asm
14{asm_x86}
15"```"
16### Response:
17```asm
18{asm_arm}
19"""
20
21# Load the model
22model = AutoModelForCausalLM.from_pretrained(
23 model_name,
24 token=hf_token,
25 device_map="auto",
26 torch_dtype=torch.bfloat16,
27)
28
29model.config.use_cache = True
30
31# Load the tokenizer
32tokenizer = AutoTokenizer.from_pretrained(
33 model_name,
34 trust_remote_code=True,
35 token=hf_token,
36)
37
38def inference(asm_x86: str) -> str:
39 prompt = instruction.format(asm_x86=asm_x86, asm_arm="")
40 inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
41 generated_ids = model.generate(
42 **inputs,
43 use_cache=True,
44 num_return_sequences=1,
45 max_new_tokens=8000,
46 do_sample=False,
47 num_beams=4,
48 eos_token_id=tokenizer.eos_token_id,
49 pad_token_id=tokenizer.pad_token_id,
50 )
51 outputs = tokenizer.batch_decode(generated_ids)[0]
52 torch.cuda.empty_cache()
53 torch.cuda.synchronize()
54 return outputs.split("```asm\n")[-1].split(f"```{tokenizer.eos_token}")[0]
55
56# Example usage
57x86 = "DWORD PTR -248[rbp] movsx rdx"
58converted_arm = inference(x86)
59print(converted_arm)| 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% |
@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}
}