A simple Phi-2 model fine-tuned on a function identification task of disassembled binary functions. It will output function names as a JSON object. You can use the following code to identify a function name:
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
model = AutoModelForCausalLM.from_pretrained(
"seanmor5/phi-2-function-identification",
attn_implementation="flash_attention_2",
torch_dtype=torch.bfloat16,
)
model.to(torch.device("cuda"))
tokenizer = AutoTokenizer.from_pretrained("seanmor5/phi-2-function-identification")
def prompt(code):
return (
"Input: Given the following disassembled code, provide a descriptive"
+ " function name for the code. Your function name should"
+ " accurately describe the purpose of the code. It should"
+ " be formatted in C style with lowercase and snakecase."
+ f" Only output the name as valid JSON, e.g. {json.dumps({'name': 'function_name'})}"
+ f"\nCode: {code}\nOutput:"
)
def identify_function(code):
eos_tokens = tokenizer.convert_tokens_to_ids(['"}', "<|endoftext|>"])
inputs = tokenizer(prompt(func), return_tensors="pt")
inputs.to(torch.device("cuda"))
outputs = model.generate(**inputs, max_new_tokens=64, eos_token_id=eos_tokens)
text = tokenizer.batch_decode(outputs[:, inputs["input_ids"].shape[1] :])[0]
return text
func = """
void fcn.140030b80(ulong param_1, ulong param_2, ulong param_3) {
ulong uVar1; uVar1 = fcn.140030ae0(param_3);
fcn.14002efc0(param_1, param_2, uVar1); return;
}
"""
print(identify_function(func))
The model tends to repeat itself excessively, so you should set the EOS token to "}
when generating.
- Downloads last month
- 8
This model does not have enough activity to be deployed to Inference API (serverless) yet. Increase its social
visibility and check back later, or deploy to Inference Endpoints (dedicated)
instead.