Spaces:
Running
Running
| import gradio as gr | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| import re | |
| # Carrega o modelo e tokenizer - adaptar conforme a documentação específica do modelo | |
| def load_model(): | |
| try: | |
| print("Carregando modelo LLM4Binary/sk2decompile-struct-6.7b...") | |
| tokenizer = AutoTokenizer.from_pretrained("LLM4Binary/sk2decompile-struct-6.7b") | |
| model = AutoModelForCausalLM.from_pretrained( | |
| "LLM4Binary/sk2decompile-struct-6.7b", | |
| torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32, | |
| device_map="auto" if torch.cuda.is_available() else None | |
| ) | |
| print("Modelo carregado com sucesso!") | |
| return model, tokenizer | |
| except Exception as e: | |
| print(f"Erro ao carregar modelo: {e}") | |
| return None, None | |
| # Função de decompilação - precisa ser adaptada conforme a API do modelo | |
| def decompile_binary(assembly_code, max_length=512, temperature=0.7): | |
| model, tokenizer = load_model() | |
| if model is None or tokenizer is None: | |
| return "Erro ao carregar o modelo. Verifique se o nome do modelo está correto." | |
| try: | |
| # Formata o input conforme esperado pelo modelo de decompilação | |
| # Este formato precisa ser adaptado baseado na documentação do modelo específico | |
| prompt = f"Decompile the following x86 assembly to C code:\n\n{assembly_code}\n\nC code:" | |
| inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024) | |
| if torch.cuda.is_available(): | |
| inputs = {k: v.cuda() for k, v in inputs.items()} | |
| with torch.no_grad(): | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=max_length, | |
| temperature=temperature, | |
| do_sample=True, | |
| top_p=0.95, | |
| pad_token_id=tokenizer.eos_token_id | |
| ) | |
| decompiled_code = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| # Extrai apenas o código C gerado (remove o prompt original) | |
| if "C code:" in decompiled_code: | |
| decompiled_code = decompiled_code.split("C code:")[1].strip() | |
| return decompiled_code | |
| except Exception as e: | |
| return f"Erro durante a decompilação: {str(e)}" | |
| # Interface Gradio | |
| with gr.Blocks(title="Binary Decompiler") as demo: | |
| gr.Markdown("# 🧠 SK²Decompile - Binary to C Code Decompiler") | |
| gr.Markdown("Decompile x86 assembly code to readable C code using LLM4Binary/sk2decompile-struct-6.7b") | |
| with gr.Row(): | |
| with gr.Column(): | |
| assembly_input = gr.Textbox( | |
| label="Assembly Code (x86)", | |
| placeholder="Digite ou cole seu código assembly aqui...", | |
| lines=15, | |
| max_lines=30 | |
| ) | |
| with gr.Row(): | |
| max_length = gr.Slider(128, 1024, value=512, step=1, label="Max Length") | |
| temperature = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="Temperature") | |
| decompile_btn = gr.Button("Decompile", variant="primary") | |
| with gr.Column(): | |
| c_output = gr.Code( | |
| label="Decompiled C Code", | |
| language="c", | |
| lines=20 | |
| ) | |
| gr.Markdown("### Exemplo de uso:") | |
| gr.Markdown("```\nmov eax, 5\nadd eax, ebx\nret\n```") | |
| # Exemplos predefinidos | |
| examples = [ | |
| ["mov eax, 5\nadd eax, ebx\nret"], | |
| ["push ebp\nmov ebp, esp\nmov eax, [ebp+8]\nadd eax, [ebp+12]\npop ebp\nret"] | |
| ] | |
| gr.Examples( | |
| examples=examples, | |
| inputs=assembly_input, | |
| outputs=c_output, | |
| fn=decompile_binary, | |
| cache_examples=True | |
| ) | |
| decompile_btn.click( | |
| fn=decompile_binary, | |
| inputs=[assembly_input, max_length, temperature], | |
| outputs=c_output | |
| ) | |
| gr.Markdown("---") | |
| gr.Markdown("🔗 **Model**: [LLM4Binary/sk2decompile-struct-6.7b](https://huggingface.co/LLM4Binary/sk2decompile-struct-6.7b)") | |
| gr.Markdown("💡 **Note**: Este é um modelo de 6.7B parâmetros especializado em decompilação binária. Para melhores resultados, forneça funções assembly completas.") | |
| if __name__ == "__main__": | |
| demo.launch() |