Update app.py

app.py

@@ -7,20 +7,26 @@ from sentence_transformers import SentenceTransformer, util
 import torch
 import gradio as gr
 
-# Inicialização do modelo BLOOM para compreensão de linguagem natural
-bloom_tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-560m")
-bloom_model = BloomForCausalLM.from_pretrained("bigscience/bloom-560m")
-bloom_model.eval()
-
-# Inicialização do modelo all-MiniLM-L6-v2 para indexação
-indexing_model = SentenceTransformer("all-MiniLM-L6-v2")
-
-# Inicialização do CodeBERT para explicações gerativas
-tokenizer = RobertaTokenizer.from_pretrained("microsoft/codebert-base")
-codebert_model = RobertaForCausalLM.from_pretrained("microsoft/codebert-base", output_hidden_states=True)
-codebert_model.eval()
-
-# Contexto global para armazenar dados do APK
+# Inicialização dos modelos com tratamento de erros
+try:
+    # BLOOM para compreensão de linguagem natural
+    bloom_tokenizer = BloomTokenizerFast.from_pretrained("bigscience/bloom-560m")
+    bloom_model = BloomForCausalLM.from_pretrained("bigscience/bloom-560m")
+    bloom_model.eval()
+
+    # Modelo de indexação
+    indexing_model = SentenceTransformer("all-MiniLM-L6-v2")
+
+    # CodeBERT para análise de código
+    tokenizer = RobertaTokenizer.from_pretrained("microsoft/codebert-base")
+    codebert_model = RobertaForCausalLM.from_pretrained("microsoft/codebert-base")
+    codebert_model.eval()
+
+except Exception as e:
+    print(f"Erro ao carregar modelos: {str(e)}")
+    raise
+
+# Contexto global
 apk_context = {"smali": {}, "java": {}, "info": ""}
 
 def check_java():
@@ -28,9 +34,9 @@ def check_java():
         result = subprocess.run(["java", "-version"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
         print("Java está disponível:", result.stderr.decode())
     except FileNotFoundError:
-        raise EnvironmentError("Java não está instalado ou não foi encontrado no PATH.")
+        raise EnvironmentError("Java não está instalado")
     except Exception as e:
-        raise EnvironmentError(f"Erro inesperado ao verificar a instalação do Java: {str(e)}")
+        raise EnvironmentError(f"Erro ao verificar Java: {str(e)}")
 
 def install_tools():
     baksmali_path = "/usr/local/bin/baksmali.jar"
@@ -39,64 +45,54 @@ def install_tools():
     check_java()
 
     if not os.path.exists(baksmali_path):
-        print("Instalando o Baksmali...")
-        subprocess.run(
-            [
-                "curl",
-                "-L",
-                "-o",
-                baksmali_path,
-                "https://bitbucket.org/JesusFreke/smali/downloads/baksmali-2.5.2.jar",
-            ],
-            check=True,
-        )
+        print("Instalando Baksmali...")
+        subprocess.run([
+            "curl", "-L", "-o", baksmali_path,
+            "https://bitbucket.org/JesusFreke/smali/downloads/baksmali-2.5.2.jar"
+        ], check=True)
 
     jadx_zip_path = "/usr/local/bin/jadx.zip"
     if not os.path.exists(jadx_path):
-        print("Instalando o JADX...")
-        subprocess.run(
-            [
-                "curl",
-                "-L",
-                "-o",
-                jadx_zip_path,
-                "https://github.com/skylot/jadx/releases/download/v1.4.7/jadx-1.4.7.zip",
-            ],
-            check=True,
-        )
+        print("Instalando JADX...")
+        subprocess.run([
+            "curl", "-L", "-o", jadx_zip_path,
+            "https://github.com/skylot/jadx/releases/download/v1.4.7/jadx-1.4.7.zip"
+        ], check=True)
         subprocess.run(["unzip", "-o", jadx_zip_path, "-d", "/usr/local/bin/jadx"], check=True)
         if os.path.exists(jadx_path):
             subprocess.run(["chmod", "+x", jadx_path], check=True)
-        else:
-            raise FileNotFoundError("Executável do JADX não encontrado no caminho esperado.")
 
 def decompile_apk(apk_file):
-    if apk_file is None:
-        return "Nenhum arquivo enviado. Por favor, envie um arquivo APK."
+    if not apk_file:
+        return "Nenhum arquivo enviado"
 
     temp_apk_path = apk_file.name
     output_dir = tempfile.mkdtemp()
     try:
+        # Smali
         smali_output = os.path.join(output_dir, "smali")
-        subprocess.run(
-            ["java", "-jar", "/usr/local/bin/baksmali.jar", "d", temp_apk_path, "-o", smali_output],
-            check=True,
-        )
+        subprocess.run([
+            "java", "-jar", "/usr/local/bin/baksmali.jar",
+            "d", temp_apk_path, "-o", smali_output
+        ], check=True)
 
+        # JADX
         java_output = os.path.join(output_dir, "java")
-        subprocess.run(
-            ["/usr/local/bin/jadx/bin/jadx", "-d", java_output, temp_apk_path],
-            check=True,
-        )
+        subprocess.run([
+            "/usr/local/bin/jadx/bin/jadx",
+            "-d", java_output, temp_apk_path
+        ], check=True)
 
+        # Coletar arquivos
         smali_files = {}
+        java_files = {}
+
         for root, _, files in os.walk(smali_output):
             for file in files:
                 if file.endswith(".smali"):
                     with open(os.path.join(root, file), "r") as f:
                         smali_files[file] = f.read()
 
-        java_files = {}
         for root, _, files in os.walk(java_output):
             for file in files:
                 if file.endswith(".java"):
@@ -106,100 +102,143 @@ def decompile_apk(apk_file):
         apk_context["smali"] = smali_files
         apk_context["java"] = java_files
 
-        return f"Decompilação bem-sucedida. Extraídos {len(smali_files)} arquivos Smali e {len(java_files)} arquivos Java."
+        return f"Decompilação concluída: {len(smali_files)} arquivos Smali, {len(java_files)} arquivos Java"
+
+    except Exception as e:
+        return f"Erro na decompilação: {str(e)}"
+
+def process_with_bloom(text):
+    try:
+        # Preparar input
+        inputs = bloom_tokenizer(
+            text,
+            return_tensors="pt",
+            max_length=512,
+            truncation=True,
+            padding=True
+        )
+
+        # Gerar resposta
+        with torch.no_grad():
+            outputs = bloom_model.generate(
+                inputs["input_ids"],
+                max_length=200,
+                num_return_sequences=1,
+                temperature=0.7,
+                pad_token_id=bloom_tokenizer.pad_token_id
+            )
+
+        # Decodificar resposta
+        processed = bloom_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return processed
 
     except Exception as e:
-        return f"Erro durante a decompilação: {str(e)}"
-
-def process_with_bloom(user_message):
-    inputs = bloom_tokenizer(user_message, return_tensors="pt", max_length=512, truncation=True)
-    with torch.no_grad():
-        outputs = bloom_model.generate(
-            inputs["input_ids"],
-            max_length=256,
-            num_return_sequences=1,
-            temperature=0.7,
-            top_p=0.9
+        print(f"Erro no processamento BLOOM: {str(e)}")
+        return text
+
+def analyze_with_codebert(code_text, query):
+    try:
+        # Preparar prompt
+        prompt = f"Query: {query}\nCódigo: {code_text[:500]}"
+        
+        # Tokenizar com padding adequado
+        inputs = tokenizer(
+            prompt,
+            return_tensors="pt",
+            max_length=512,
+            truncation=True,
+            padding="max_length"
         )
-        processed_query = bloom_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return processed_query
 
-def get_embeddings(text):
-    return indexing_model.encode(text, convert_to_tensor=True)
+        # Gerar análise
+        with torch.no_grad():
+            outputs = codebert_model.generate(
+                inputs["input_ids"],
+                max_length=200,
+                num_return_sequences=1,
+                pad_token_id=tokenizer.pad_token_id
+            )
+
+        # Decodificar resposta
+        analysis = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return analysis
+
+    except Exception as e:
+        print(f"Erro na análise CodeBERT: {str(e)}")
+        return "Não foi possível analisar o código"
 
 def query_apk_chat(user_message):
     if not apk_context["smali"] and not apk_context["java"]:
-        return "Nenhum APK decompilado disponível. Por favor, envie e decompile um APK primeiro."
+        return "Nenhum APK decompilado disponível"
 
     try:
-        # Processar a mensagem do usuário com BLOOM
-        processed_message = process_with_bloom(user_message)
-        
-        # Obter embedding da query processada
-        query_embedding = get_embeddings(processed_message)
+        # Processar query com BLOOM
+        processed_query = process_with_bloom(user_message)
+
+        # Preparar todos os códigos
+        all_codes = []
+        all_files = []
 
-        # Combinar Smali e Java para análise
-        combined_texts = [(k, v) for k, v in apk_context["smali"].items()] + [(k, v) for k, v in apk_context["java"].items()]
-        combined_embeddings = indexing_model.encode([v for _, v in combined_texts], convert_to_tensor=True)
+        for file, code in apk_context["java"].items():
+            all_codes.append(code)
+            all_files.append(("java", file))
 
-        # Encontrar códigos relevantes
-        scores = util.pytorch_cos_sim(query_embedding, combined_embeddings).squeeze(0)
-        top_k = min(3, len(combined_texts))
-        top_indices = torch.topk(scores, k=top_k).indices
+        for file, code in apk_context["smali"].items():
+            all_codes.append(code)
+            all_files.append(("smali", file))
 
-        response = ""
-        for idx in top_indices:
-            file_name, relevant_code = combined_texts[idx.item()]
+        # Calcular embeddings
+        query_embedding = indexing_model.encode(processed_query, convert_to_tensor=True)
+        code_embeddings = indexing_model.encode(all_codes, convert_to_tensor=True)
+
+        # Encontrar matches
+        similarities = util.pytorch_cos_sim(query_embedding, code_embeddings)[0]
+        top_k = min(3, len(all_codes))
+        best_matches = torch.topk(similarities, k=top_k)
+        
+        response = []
+        for score, idx in zip(best_matches.values, best_matches.indices):
+            file_type, file_name = all_files[idx]
+            code = all_codes[idx]
             
-            # Gerar explicação usando CodeBERT com a query processada
-            explanation_prompt = f"Query processada: {processed_message}\nExplique o código:\n{relevant_code[:500]}"
-            inputs = tokenizer(explanation_prompt, return_tensors="pt", max_length=512, truncation=True, padding=True)
+            # Análise do código
+            analysis = analyze_with_codebert(code, processed_query)
             
-            with torch.no_grad():
-                outputs = codebert_model.generate(
-                    inputs["input_ids"],
-                    max_length=1024,
-                    num_return_sequences=1,
-                    temperature=0.7,
-                    top_p=0.9
-                )
-                explanation = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-            response += f"\n\n**Arquivo:** {file_name}\n"
-            response += f"**Código:**\n```\n{relevant_code[:1000]}\n```\n"
-            response += f"**Explicação:**\n{explanation}\n"
-            response += "-" * 80
-
-        return response
+            response.append(f"\nArquivo ({file_type}): {file_name}")
+            response.append(f"Relevância: {score:.2f}")
+            response.append(f"Código:\n{code[:500]}...")
+            response.append(f"Análise:\n{analysis}\n")
+            response.append("-" * 80)
+
+        return "\n".join(response)
 
     except Exception as e:
-        return f"Erro durante a análise: {str(e)}"
+        return f"Erro na análise: {str(e)}"
 
-# Configuração e inicialização
+# Configuração Gradio
 install_tools()
 
-# Interfaces Gradio
-apk_upload_interface = gr.Interface(
+upload_interface = gr.Interface(
     fn=decompile_apk,
-    inputs=gr.File(label="Enviar arquivo APK", file_types=[".apk"]),
+    inputs=gr.File(label="APK File", file_types=[".apk"]),
     outputs="text",
-    title="Analisador de APK com BLOOM + CodeBERT",
-    description="Envie um arquivo APK para análise avançada com IA."
+    title="APK Analyzer",
+    description="Upload an APK file for analysis"
 )
 
 chat_interface = gr.Interface(
     fn=query_apk_chat,
-    inputs=gr.Textbox(lines=3, placeholder="Faça uma pergunta sobre o código do APK..."),
-    outputs=gr.Textbox(lines=20, label="Análise Detalhada"),
-    title="Chat Avançado com APK",
-    description="Análise inteligente do código usando BLOOM para processamento de linguagem natural e CodeBERT para análise técnica."
+    inputs=gr.Textbox(lines=3, placeholder="Ask about the APK code..."),
+    outputs=gr.Textbox(lines=20),
+    title="Code Analysis Chat",
+    description="AI-powered code analysis"
 )
 
 # Interface combinada
 iface = gr.TabbedInterface(
-    [apk_upload_interface, chat_interface],
-    ["Enviar & Analisar", "Análise Inteligente"]
+    [upload_interface, chat_interface],
+    ["Upload APK", "Analyze Code"]
 )
 
-# Iniciar a interface
-iface.launch()
+if __name__ == "__main__":
+    iface.launch()