VirtualCompiler / do_request.py
elsagranger's picture
Add usage
78bb3c5
import requests
import random
CL = ['clang-11', 'clang-12', 'clang-9', 'gcc-11', 'gcc-7', 'gcc-9']
OP = ['O0', 'O1', 'O2', 'O3', 'Os']
ST = ['stripped', 'unstripped']
def process(source_code):
compiler = random.choice(CL)
optimizer = random.choice(OP)
strip_type = random.choice(ST)
prompt = f'Please compile this source code using {compiler} with optimization level {optimizer} into assembly code.'
if strip_type == 'stripped':
prompt += ' Strip the assembly code.'
else:
prompt += ' No strip the assembly code.'
query_prompt = "<s>system\n" + prompt + \
"</s>\n<s>user\n" + source_code + "</s>\n<s>assistant\n"
return query_prompt, compiler, optimizer, strip_type
def do_request(src):
url = "http://localhost:8080/v1/completions"
query_prompt, _, _, _ = process(src)
model_name = "VirtualCompiler"
ret = requests.post(url, json={
"prompt": query_prompt,
"max_tokens": 4096,
"temperature": 0.3,
"stop": ["</s>"],
"model": model_name,
"echo": False,
"logprobs": True,
})
return ret.json()
src = '''static int
layout_append(struct layout_cell *lc, char *buf, size_t len)
{
if (len == 0)
return (-1);
return (0);
}
'''
ret = do_request(src)
print(ret['choices'][0]['text'])