Virtualcompiler

🚀 虛擬編譯器：彙編代碼搜索的理想之選

本項目為ACL 2024論文 “Virtual Compiler Is All You Need For Assembly Code Search” 提供了模型和相應的評估數據集。虛擬編譯器是一種大語言模型（LLM），它能夠將任何編程語言編譯為底層的彙編代碼。該虛擬編譯器模型基於34B CodeLlama，可在 elsagranger/VirtualCompiler 獲取。

🚀 快速開始

本項目包含模型和評估數據集，下面為你介紹使用方法。

✨ 主要特性

• 虛擬編譯器模型：基於34B CodeLlama的大語言模型，可將任意編程語言編譯為彙編代碼。
• 多方面評估：通過腳本強制執行評估虛擬彙編代碼與真實彙編代碼的相似度，還通過彙編代碼搜索這一下游任務評估虛擬編譯器的有效性。

📦 安裝指南

我們使用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"]

# Extracted by the process_asm.py script mentioned above
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])

📚 詳細文檔

• 相似度評估：使用腳本 force-exec.py 強制執行，評估虛擬編譯器生成的虛擬彙編代碼與真實彙編代碼的相似度。相應的評估數據集可在 virtual_assembly_and_ground_truth 獲取。
• 有效性評估：通過彙編代碼搜索這一下游任務評估虛擬編譯器的有效性。評估數據集可在 elsagranger/AssemblyCodeSearchEval 獲取。

📄 許可證

本項目採用 Apache-2.0 許可證。