🚀 アセンブリコード検索に必要なすべてがVirtual Compiler
このリポジトリには、ACL 2024の論文 "Virtual Compiler Is All You Need For Assembly Code Search" のモデルと対応する評価データセットが含まれています。
仮想コンパイラは、あらゆるプログラミング言語を基礎となるアセンブリコードにコンパイルできる大規模言語モデル(LLM)です。仮想コンパイラモデルは、elsagranger/VirtualCompiler で利用可能で、34BのCodeLlamaをベースにしています。
私たちは、仮想コンパイラによって生成された仮想アセンブリコードと実際のアセンブリコードの類似性を、スクリプト force-exec.py による強制実行を使用して評価しています。対応する評価データセットは、virtual_assembly_and_ground_truth で利用可能です。
また、仮想コンパイラの有効性を下流タスクであるアセンブリコード検索を通じて評価しています。評価データセットは、elsagranger/AssemblyCodeSearchEval で利用可能です。
🚀 クイックスタート
💻 使用例
基本的な使用法
私たちはFastChatとvllm workerを使用してモデルをホストします。以下のコマンドを別々のターミナル(例えば tmux
)で実行してください。
LOGDIR="" python3 -m fastchat.serve.openai_api_server \
--host 0.0.0.0 --port 8080 \
--controller-address http://localhost:21000
LOGDIR="" python3 -m fastchat.serve.controller \
--host 0.0.0.0 --port 21000
LOGDIR="" RAY_LOG_TO_STDERR=1 \
python3 -m fastchat.serve.vllm_worker \
--model-path ./VirtualCompiler \
--num-gpus 8 \
--controller http://localhost:21000 \
--max-num-batched-tokens 40960 \
--disable-log-requests \
--host 0.0.0.0 --port 22000 \
--worker-address http://localhost:22000 \
--model-names "VirtualCompiler"
モデルがホストされたら、do_request.py
を使用してモデルにリクエストを送信します。
~/C/VirtualCompiler (main)> python3 do_request.py
test rdx, rdx
setz al
movzx eax, al
neg eax
retn
高度な使用法
Hugging Faceはフォルダ内のリモートモデルの読み込みをサポートしていないため、私たちは仮想コンパイラによって拡張されたアセンブリコード検索データセットで学習されたモデルを vic-encoder でホストしています。カスタムモデルの読み込みをテストするには、model.py
を使用できます。
以下は、テキストエンコーダとアセンブリエンコーダの使用例です。バイナリからアセンブリコードを抽出する方法については、このスクリプト process_asm.py を参照してください。
def calc_map_at_k(logits, pos_cnt, ks=[10,]):
_, indices = torch.sort(logits, dim=1, descending=True)
ranks = torch.nonzero(
indices < pos_cnt,
as_tuple=False
)[:, 1].reshape(logits.shape[0], -1)
mrr = torch.mean(1 / (ranks + 1), dim=1)
res = {}
for k in ks:
res[k] = (
torch.sum((ranks < k).float(), dim=1) / min(k, pos_cnt)
).cpu().numpy()
return ranks.cpu().numpy(), res, mrr.cpu().numpy()
pos_asm_cnt = 1
query = ["List all files in a directory"]
anchor_asm = [ {"1": "endbr64", "2": "mov eax, 0" }, ... ]
neg_anchor_asm = [ {"1": "push rbp", "2": "mov rbp, rsp", ... }, ... ]
query_embs = text_encoder(**text_tokenizer(query))
kwargs = dict(padding=True, pad_to_multiple_of=8, return_tensors="pt")
anchor_asm_ids = asm_tokenizer.pad([asm_tokenizer(pos) for pos in anchor_asm], **kwargs)
neg_anchor_asm_ids = asm_tokenizer.pad([asm_tokenizer(neg) for neg in neg_anchor_asm], **kwargs)
asm_embs = asm_encoder(**anchor_asm_ids)
asm_neg_emb = asm_encoder(**neg_anchor_asm_ids)
logits_pos = torch.einsum(
"ic,jc->ij", [query_embs, asm_embs])
logits_neg = torch.einsum(
"ic,jc->ij", [query_embs, asm_neg_emb[pos_asm_cnt:]]
)
logits = torch.cat([logits_pos, logits_neg], dim=1)
ranks, map_at_k, mrr = calc_map_at_k(
logits, pos_asm_cnt, [1, 5, 10, 20, 50, 100])
📄 ライセンス
このプロジェクトは、Apache-2.0ライセンスの下で公開されています。