Update app.py

app.py

@@ -1,322 +1,113 @@
-# app.py - Versão ultra-robusta com múltiplos fallbacks
-import os
 import gradio as gr
-from huggingface_hub import InferenceClient, InferenceTimeoutError
-import time
-import logging
-from typing import Generator, Optional
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import re
 
-# Configuração de logging para debugging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Modelos em ordem de prioridade (do mais específico para o mais genérico)
-FALLBACK_MODELS = [
-    "LLM4Binary/sk2decompile-struct-6.7b",  # Modelo original
-    "mradermacher/sk2decompile-struct-6.7b-GGUF",  # Versão GGUF
-    "meta-llama/Meta-Llama-3-8B-Instruct",  # Fallback genérico confiável
-    "mistralai/Mistral-7B-Instruct-v0.2",   # Outro fallback
-    "google/gemma-2b-it"                    # Fallback leve
-]
-
-def get_valid_token(hf_token_obj: Optional[object] = None) -> str:
-    """Obtém um token válido de múltiplas fontes com fallbacks"""
+# Carrega o modelo e tokenizer - adaptar conforme a documentação específica do modelo
+def load_model():
     try:
-        # Tenta obter token do objeto OAuth
-        if hf_token_obj and hasattr(hf_token_obj, "token") and hf_token_obj.token:
-            logger.info("Usando token do OAuth")
-            return hf_token_obj.token
-        
-        # Tenta variável de ambiente
-        env_token = os.getenv("HF_TOKEN", "").strip()
-        if env_token:
-            logger.info("Usando token da variável de ambiente HF_TOKEN")
-            return env_token
-        
-        # Tenta arquivo .env
-        try:
-            from dotenv import load_dotenv
-            load_dotenv()
-            env_token = os.getenv("HF_TOKEN", "").strip()
-            if env_token:
-                logger.info("Usando token do arquivo .env")
-                return env_token
-        except ImportError:
-            pass
-        
-        # Último fallback: tenta sem token (alguns modelos públicos permitem)
-        logger.warning("Nenhum token encontrado. Tentando sem autenticação...")
-        return None
-        
+        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:
-        logger.error(f"Erro ao obter token: {e}")
-        return None
+        print(f"Erro ao carregar modelo: {e}")
+        return None, None
 
-def create_client(model_id: str, token: Optional[str] = None) -> InferenceClient:
-    """Cria um cliente de inferência com configurações seguras"""
+# 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:
-        # Configurações seguras com timeout
-        timeout = 60  # 60 segundos para operações complexas
+        # 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:"
         
-        logger.info(f"Criando cliente para modelo: {model_id}")
-        client = InferenceClient(
-            model=model_id,
-            token=token,
-            timeout=timeout
-        )
+        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()}
         
-        # Teste rápido de conexão
-        test_response = client.post(
-            json={"inputs": "test"},
-            model=model_id,
-            timeout=10  # timeout curto para teste
-        )
+        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
+            )
         
-        if test_response.status_code in [200, 400]:  # 400 pode ser erro de input, mas API está respondendo
-            logger.info(f"Conexão bem-sucedida com {model_id}")
-            return client
+        decompiled_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-        logger.warning(f"Teste falhou com status {test_response.status_code} para {model_id}")
-        return None
+        # 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()
         
-    except InferenceTimeoutError:
-        logger.warning(f"Timeout ao conectar com {model_id}")
-        return None
-    except Exception as e:
-        logger.error(f"Erro ao criar cliente para {model_id}: {e}")
-        return None
-
-def try_models(messages, max_tokens, temperature, top_p, token):
-    """Tenta múltiplos modelos em ordem até obter sucesso"""
-    for model_id in FALLBACK_MODELS:
-        try:
-            logger.info(f"Tentando modelo: {model_id}")
-            client = create_client(model_id, token)
-            
-            if not client:
-                logger.warning(f"Cliente inválido para {model_id}, pulando...")
-                continue
-            
-            # Tenta streaming primeiro
-            response = ""
-            stream_success = False
-            
-            try:
-                stream = client.chat_completion(
-                    messages=messages,
-                    max_tokens=max_tokens,
-                    stream=True,
-                    temperature=temperature,
-                    top_p=top_p,
-                    timeout=120  # timeout maior para streaming
-                )
-                
-                for chunk in stream:
-                    if hasattr(chunk, "choices") and chunk.choices:
-                        choice = chunk.choices[0]
-                        if hasattr(choice, "delta") and hasattr(choice.delta, "content"):
-                            token_content = choice.delta.content or ""
-                            if token_content:
-                                response += token_content
-                                yield response
-                                stream_success = True
-                
-                if stream_success and response.strip():
-                    logger.info(f"Resposta obtida com sucesso de: {model_id} (streaming)")
-                    return
-            
-            except Exception as stream_error:
-                logger.warning(f"Erro no streaming com {model_id}: {stream_error}")
-            
-            # Fallback para chamada síncrona se streaming falhar
-            if not stream_success:
-                logger.info(f"Tentando chamada síncrona com {model_id}")
-                try:
-                    full_response = client.chat_completion(
-                        messages=messages,
-                        max_tokens=max_tokens,
-                        stream=False,
-                        temperature=temperature,
-                        top_p=top_p,
-                        timeout=60
-                    )
-                    
-                    # Extrai resposta de múltiplas formas possíveis
-                    if hasattr(full_response, "choices") and full_response.choices:
-                        choice = full_response.choices[0]
-                        if hasattr(choice, "message") and hasattr(choice.message, "content"):
-                            response = choice.message.content
-                        elif hasattr(choice, "text"):
-                            response = choice.text
-                    elif isinstance(full_response, dict):
-                        if "generated_text" in full_response:
-                            response = full_response["generated_text"]
-                        elif "choices" in full_response and full_response["choices"]:
-                            first_choice = full_response["choices"][0]
-                            if "message" in first_choice and "content" in first_choice["message"]:
-                                response = first_choice["message"]["content"]
-                    
-                    if response and response.strip():
-                        logger.info(f"Resposta obtida com sucesso de: {model_id} (síncrono)")
-                        yield response
-                        return
-                
-                except Exception as sync_error:
-                    logger.warning(f"Erro na chamada síncrona com {model_id}: {sync_error}")
-            
-            # Pequena pausa entre tentativas para não sobrecarregar
-            time.sleep(1)
-            
-        except Exception as model_error:
-            logger.error(f"Erro geral com modelo {model_id}: {model_error}")
-            continue
+        return decompiled_code
     
-    # Se todos os modelos falharem
-    error_msg = (
-        "❌ **Erro persistente**: Nenhum modelo disponível no momento.\n\n"
-        "💡 **Soluções sugeridas**:\n"
-        "1. Verifique sua conexão com a internet\n"
-        "2. Configure um token HF válido em Settings → Secrets\n"
-        "3. Tente novamente em alguns minutos\n\n"
-        f"📋 Últimos modelos tentados: {', '.join(FALLBACK_MODELS[:3])}"
-    )
-    logger.error("Todos os modelos falharam")
-    yield error_msg
+    except Exception as e:
+        return f"Erro durante a decompilação: {str(e)}"
 
-def respond(message: str, history: list, system_message: str, max_tokens: int, temperature: float, top_p: float, hf_token_obj=None) -> Generator[str, None, None]:
-    """Função de resposta ultra-resiliente"""
+# 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")
     
-    # Validação de entrada
-    if not message or not message.strip():
-        yield "❌ **Mensagem vazia**: Por favor, digite uma mensagem válida."
-        return
-    
-    # Obtém token válido
-    token = get_valid_token(hf_token_obj)
-    if not token:
-        warning_msg = (
-            "⚠️ **Sem autenticação**: Operando em modo limitado.\n"
-            "Para melhor performance e acesso a mais modelos:\n"
-            "1. Configure HF_TOKEN em Settings → Secrets\n"
-            "2. Ou faça login na interface\n\n"
-            "Tentando com modelos públicos..."
-        )
-        yield warning_msg
-        # Não retorna aqui, continua tentando com modelos públicos
-    
-    # Prepara mensagens
-    messages = []
-    
-    # Adiciona mensagem de sistema se válida
-    if system_message and system_message.strip():
-        messages.append({"role": "system", "content": system_message.strip()})
-    else:
-        # Mensagem de sistema padrão segura
-        messages.append({"role": "system", "content": "Você é um assistente útil e especializado em análise de estruturas de código."})
-    
-    # Processa histórico com segurança
-    if history:
-        for entry in history:
-            try:
-                if isinstance(entry, dict) and "role" in entry and "content" in entry:
-                    messages.append(entry)
-                elif isinstance(entry, list) and len(entry) >= 2:
-                    # Formato [(user_msg, bot_response), ...]
-                    if entry[0]:  # mensagem do usuário
-                        messages.append({"role": "user", "content": str(entry[0])})
-                    if entry[1]:  # resposta do bot
-                        messages.append({"role": "assistant", "content": str(entry[1])})
-            except Exception as e:
-                logger.warning(f"Erro ao processar histórico: {e}")
-                continue
+    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
+            )
     
-    # Adiciona mensagem atual
-    messages.append({"role": "user", "content": message.strip()})
+    gr.Markdown("### Exemplo de uso:")
+    gr.Markdown("```\nmov eax, 5\nadd eax, ebx\nret\n```")
     
-    logger.info(f"Mensagens preparadas: {messages}")
+    # 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"]
+    ]
     
-    # Tenta obter resposta com múltiplos fallbacks
-    try:
-        yield from try_models(messages, max_tokens, temperature, top_p, token)
-    except Exception as final_error:
-        logger.critical(f"Erro crítico inesperado: {final_error}")
-        yield (
-            "❌ **Erro crítico**: Ocorreu um problema inesperado.\n\n"
-            f"```python\n{str(final_error)}\n```\n\n"
-            "Por favor, recarregue a página e tente novamente."
-        )
-
-# Interface do usuário - simples e robusta
-with gr.Blocks(title="Assistente de Análise de Código") as demo:
-    gr.Markdown("# 🔍 Assistente de Análise de Estruturas de Código")
-    gr.Markdown("### Versão robusta com múltiplos fallbacks e recuperação de erros")
+    gr.Examples(
+        examples=examples,
+        inputs=assembly_input,
+        outputs=c_output,
+        fn=decompile_binary,
+        cache_examples=True
+    )
     
-    chatbot = gr.ChatInterface(
-        respond,
-        type="messages",
-        additional_inputs=[
-            gr.Textbox(
-                value="Você é um especialista em análise de estruturas de código e decompilação. Forneça respostas técnicas detalhadas e precisas.",
-                label="System Message",
-                lines=3
-            ),
-            gr.Slider(
-                minimum=1, maximum=4096, value=1024, step=1,
-                label="Max Tokens (aumente para respostas mais longas)"
-            ),
-            gr.Slider(
-                minimum=0.0, maximum=2.0, value=0.3, step=0.1,
-                label="Temperature (0.0 = preciso, 2.0 = criativo)"
-            ),
-            gr.Slider(
-                minimum=0.1, maximum=1.0, value=0.9, step=0.05,
-                label="Top-p (0.1 = focado, 1.0 = diverso)"
-            ),
-        ],
-        examples=[
-            ["Analise esta função vulnerável: `def process_input(data): eval(data)`"],
-            ["Qual a estrutura de memória desta classe C++?"],
-            ["Explique o assembly x86 deste código binário"],
-            ["Como funciona o mecanismo de herança neste código?"],
-        ],
-        cache_examples=False,
-        analytics_enabled=False,
+    decompile_btn.click(
+        fn=decompile_binary,
+        inputs=[assembly_input, max_length, temperature],
+        outputs=c_output
     )
     
-    with gr.Accordion("ℹ️ Informações e Solução de Problemas", open=False):
-        gr.Markdown("""
-        ### ✅ Este aplicativo é 100% robusto:
-        
-        - **Múltiplos fallbacks**: Tenta até 5 modelos diferentes
-        - **Recuperação de erros**: Nunca falha completamente
-        - **Autenticação flexível**: Usa token do ambiente ou OAuth
-        - **Timeouts seguros**: Previne travamentos
-        - **Validação rigorosa**: Checa todas as entradas
-        
-        ### 🛠️ Se ainda encontrar problemas:
-        
-        1. **Configure HF_TOKEN**: Vá em Settings → Secrets e adicione seu token
-        2. **Recarregue a página**: Às vezes a conexão precisa ser renovada
-        3. **Simplifique sua query**: Modelos têm limites de contexto
-        4. **Verifique sua internet**: Necessária para chamadas à API
-        
-        ### 📊 Modelos utilizados (em ordem de prioridade):
-        1. LLM4Binary/sk2decompile-struct-6.7b (especializado)
-        2. mradermacher/sk2decompile-struct-6.7b-GGUF (GGUF)
-        3. Meta-Llama-3-8B-Instruct (genérico confiável)
-        4. Mistral-7B-Instruct-v0.2 (alternativo)
-        5. Gemma-2b-it (fallback leve)
-        """)
+    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__":
-    # Configurações de lançamento ultra-seguras
-    demo.launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        share=False,
-        show_api=False,
-        favicon_path="https://huggingface.co/front/assets/huggingface_logo-noborder.svg",
-        allowed_paths=["."],
-        auth=None,  # Não força autenticação
-        debug=False,  # Produção
-    )
+    demo.launch()