Upload 7 files

README.md

@@ -1,4 +1,91 @@
 ---
-tags:
-- code
----
+license: mit
+---
+
+<h1 align="center">OpTrans: Enhancing Binary Code Similarity Detection with Function Inlining Re-Optimization</h1>
+
+<h4 align="center">
+<p>
+<a href=#about>About</a> |
+<a href=#intuition>Intuition</a> |
+<a href=#quickstart>QuickStart</a> |
+<p>
+</h4>
+
+## About
+
+OpTrans (Re-Optimization Transformer), is an innovative framework fuses binary code optimization techniques with the transformer model for BCSD. By OpTrans employs an algorithm based on binary program analysis to determine which functions should be inlined, followed by binary rewriting techniques to effectuate re-optimization on binaries. Our goal is to provide an effective tool for researchers and practitioners in binary code similarity detection, with our models accessible on the Hugging Face Model Hub.
+
+
+## Intuition
+
+This document will present how function inlining optimization improve binary code similarity detection.
+
+Function Faust_next in sc3-plugins-HOAEncLebedev501.so compiled with -O0 (sc3-plugins-HOAEncLebedev501.so-O0.i64)
+![O0](./Intuition/O0.jpg)
+
+Function Faust_next in sc3-plugins-HOAEncLebedev501.so compiled with -O3 (sc3-plugins-HOAEncLebedev501.so-O3.i64)
+![O3](./Intuition/O3.jpg)
+
+Function Faust_next in sc3-plugins-HOAEncLebedev501.so compiled with -O0 and processed by function inlining optimization (sc3-plugins-HOAEncLebedev501.so-O0-inline.i64)
+![O0-inline](./Intuition/O0-inline.jpg)
+
+
+## QuickStart
+
+This document will help you set up and start using the OpTrans model for embedding generation.
+
+
+### Requirements
+- Python 3.6 or higher
+- [PyTorch](https://pytorch.org/get-started/locally/)
+- [Transformers library](https://huggingface.co/docs/transformers/installation)
+- A CUDA-enabled GPU is highly recommended for faster processing.
+
+Ensure you have Python and PyTorch installed on your system. Then, install the Transformers library using pip:
+```bash
+pip install transformers
+```
+
+### Preparing Tokenizers and Models
+Import necessary libraries and initialize the model and tokenizers:
+```python
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+tokenizer = AutoTokenizer.from_pretrained("sandspeare/optrans", trust_remote_code=True)
+encoder = AutoModel.from_pretrained("sandspeare/optrans", trust_remote_code=True).to(device)
+tokenizer.pad_token = tokenizer.unk_token
+```
+
+### Example Use Cases
+**Function inlining optimization for BCSD**
+
+1. Load your binary code dataset. For demonstration, we use a pickle file containing binary code snippets for similarity compare.
+
+```python
+with open("./CaseStudy/casestudy.json") as fp:
+   data = json.load(fp)
+```
+
+2. Encode the binary code.
+
+```python
+asm_O0 = tokenizer([data["O0"]], padding=True, return_tensors="pt").to(device)
+asm_embedding_O0 = encoder(**asm_O0)
+
+asm_O0_inline = tokenizer([data["O0_inline"]], padding=True, return_tensors="pt").to(device)
+asm_embedding_O0_inline = encoder(**asm_O0_inline)
+
+asm_O3 = tokenizer([data["O3"]], padding=True, return_tensors="pt").to(device)
+asm_embedding_O3 = encoder(**asm_O3)
+```
+
+3. Perform similarity comparison:
+
+```python
+sim_O0vsO3 = torch.mm(asm_embedding_O0, asm_embedding_O3.T).squeeze() / 0.07
+sim_O0_inlinevsO3 = torch.mm(asm_embedding_O0_inline, asm_embedding_O3.T).squeeze() / 0.07
+```

config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "AsmEncoder"
+  ],
+  "auto_map": {
+    "AutoModel": "optrans_modeling.AsmEncoder"
+  },
+  "attention_probs_dropout_prob": 0.1,
+  "bla_dim": 1024,
+  "embedding_size": 768,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 2048,
+  "model_type": "roformer",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 1,
+  "rotary_value": false,
+  "token_type_ids": "share",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.35.0",
+  "type_vocab_size": 2,
+  "use_bias": false,
+  "use_cache": true,
+  "vocab_size": 33494
+}

mytrans_modeling.py

@@ -0,0 +1,138 @@
+import numpy as np
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from typing import Optional, Tuple
+import torch.nn.functional as F
+from transformers import BatchEncoding
+from transformers import MPNetTokenizerFast
+
+
+from transformers.models.roformer.modeling_roformer import (
+    RoFormerEmbeddings,
+    RoFormerModel,
+    RoFormerEncoder,
+    RoFormerLayer,
+    RoFormerAttention,
+    RoFormerIntermediate,
+    RoFormerOutput,
+    RoFormerSelfAttention,
+    RoFormerPreTrainedModel
+)
+
+from transformers.models.mpnet.modeling_mpnet import MPNetModel
+
+
+class JRoFormerEmbeddings(RoFormerEmbeddings):
+    """Construct the embeddings from word and token_type embeddings."""
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.word_embeddings = nn.Embedding(
+            config.vocab_size, config.embedding_size, padding_idx=config.pad_token_id
+        )
+        self.token_type_embeddings = self.word_embeddings
+
+
+class JRoFormerSelfAttention(RoFormerSelfAttention):
+    def __init__(self, config):
+        super().__init__(config)
+        self.query = nn.Linear(
+            config.hidden_size, self.all_head_size, bias=config.use_bias
+        )
+        self.key = nn.Linear(
+            config.hidden_size, self.all_head_size, bias=config.use_bias
+        )
+        self.value = nn.Linear(
+            config.hidden_size, self.all_head_size, bias=config.use_bias
+        )
+
+
+class JRoFormerAttention(RoFormerAttention):
+    def __init__(self, config):
+        super().__init__(config)
+        self.self = JRoFormerSelfAttention(config)
+
+
+class JRoFormerLayer(RoFormerLayer):
+    def __init__(self, config):
+        super().__init__(config)
+        self.attention = JRoFormerAttention(config)
+        self.is_decoder = config.is_decoder
+        self.add_cross_attention = config.add_cross_attention
+        if self.add_cross_attention:
+            if not self.is_decoder:
+                raise ValueError(
+                    f"{self} should be used as a decoder model if cross attention is added"
+                )
+            self.crossattention = RoFormerAttention(config)
+        self.intermediate = RoFormerIntermediate(config)
+        self.output = RoFormerOutput(config)
+
+
+class JRoFormerEncoder(RoFormerEncoder):
+    def __init__(self, config):
+        super().__init__(config)
+        self.layer = nn.ModuleList(
+            [JRoFormerLayer(config) for _ in range(config.num_hidden_layers)]
+        )
+
+
+class JRoFormerModel(RoFormerModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.embeddings = JRoFormerEmbeddings(config)
+
+        if config.embedding_size != config.hidden_size:
+            self.embeddings_project = nn.Linear(
+                config.embedding_size, config.hidden_size
+            )
+
+        self.encoder = JRoFormerEncoder(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+class AsmEncoder(RoFormerPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.roformer = JRoFormerModel(config)
+        self.projection = nn.Linear(config.hidden_size, config.bla_dim)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        token_type_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        encoder_hidden_states: Optional[torch.FloatTensor] = None,
+        encoder_attention_mask: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.roformer(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        token_embeddings = outputs[0]
+        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).to(token_embeddings.dtype)
+        asm_embedding = torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+        asm_embedding = self.projection(asm_embedding)
+        asm_embedding = F.normalize(asm_embedding, p=2, dim=1)
+
+        return asm_embedding

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4add8e8eb42e5e0a6e6864837fb0da7b6fe547a2b171b99046efd09478e65a05
+size 223364289

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}