| import os | |
| import sys | |
| from llama_cpp import Llama | |
| from util import create_prompt, is_bytecode_empty | |
| llm = Llama( | |
| model_path="sentient-simulations-pydecompiler-3.7-6.7b-v0.9-q8_0.gguf", | |
| n_gpu_layers=-1, | |
| last_n_tokens_size=0, | |
| n_ctx=16384, | |
| ) | |
| def decompile_pycb(file_path: str, force_overwrite=False): | |
| source_code_path = file_path[:-2] | |
| with open(file_path, "r") as bytecode_file: | |
| bytecode = bytecode_file.read() | |
| if is_bytecode_empty(bytecode): | |
| return "# Empty file" | |
| prompt = create_prompt(bytecode) | |
| tokens = llm.tokenize(text=prompt.encode("utf-8"), add_bos=True, special=True) | |
| if len(tokens) > 15000: | |
| return None | |
| output = llm(prompt, max_tokens=None, echo=False) | |
| result = output["choices"][0] | |
| source_code = result.get("text").strip() | |
| finish_reason = result.get("finish_reason") | |
| source_code += f"\n\n# Finish Reason: {finish_reason}\n" | |
| with open(source_code_path, "w") as source_code_file: | |
| source_code_file.write(source_code) | |
| return source_code | |
| if __name__ == "__main__": | |
| path = sys.argv[1] | |
| if os.path.isfile(path): | |
| output = decompile_pycb(path) | |
| print(output) | |
| elif os.path.isdir(path): | |
| for root, dirs, files in os.walk(path): | |
| for file in files: | |
| if file.endswith(".pycb") and not os.path.exists( | |
| os.path.join(root, file[:-2]) | |
| ): | |
| file_path = os.path.join(root, file) | |
| output = decompile_pycb(file_path) | |
| if output: | |
| print(f"Decompiled {file_path}") | |