Virtualcompiler

🚀 アセンブリコード検索に必要なすべてが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)

    # [batch_size, pos_cnt]
    ranks = torch.nonzero(
        indices < pos_cnt,
        as_tuple=False
    )[:, 1].reshape(logits.shape[0], -1)

    # [batch_size, pos_cnt]
    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"]

# 上記のprocess_asm.pyスクリプトによって抽出されたもの
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)

# query_embs: [query_cnt, emb_dim]
# asm_embs: [pos_asm_cnt, emb_dim]

# logits_pos: [query_cnt, pos_asm_cnt]
logits_pos = torch.einsum(
    "ic,jc->ij", [query_embs, asm_embs])
# logits_neg: [query_cnt, neg_asm_cnt]
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ライセンスの下で公開されています。