Update app.py

app.py

@@ -2,27 +2,27 @@ import os
 import subprocess
 import tempfile
 from androguard.misc import AnalyzeAPK
-from transformers import RobertaTokenizer, RobertaForCausalLM, pipeline
+from transformers import BloomForCausalLM, BloomTokenizerFast, RobertaTokenizer, RobertaForCausalLM, pipeline
 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
-try:
-    tokenizer = RobertaTokenizer.from_pretrained("microsoft/codebert-base")
-    codebert_model = RobertaForCausalLM.from_pretrained("microsoft/codebert-base", output_hidden_states=True)
-    codebert_model.eval()  # Colocar o modelo em modo de avaliação
-except Exception as e:
-    print(f"Erro ao carregar o modelo CodeBERT: {str(e)}")
+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
 apk_context = {"smali": {}, "java": {}, "info": ""}
 
-
-# Função para verificar se o Java está instalado
 def check_java():
     try:
         result = subprocess.run(["java", "-version"], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
@@ -32,16 +32,12 @@ def check_java():
     except Exception as e:
         raise EnvironmentError(f"Erro inesperado ao verificar a instalação do Java: {str(e)}")
 
-
-# Instalar ferramentas como Baksmali e JADX
 def install_tools():
     baksmali_path = "/usr/local/bin/baksmali.jar"
     jadx_path = "/usr/local/bin/jadx/bin/jadx"
 
-    # Verificar se o Java está disponível
     check_java()
 
-    # Instalar o Baksmali caso não esteja presente
     if not os.path.exists(baksmali_path):
         print("Instalando o Baksmali...")
         subprocess.run(
@@ -55,7 +51,6 @@ def install_tools():
             check=True,
         )
 
-    # Instalar o JADX caso não esteja presente
     jadx_zip_path = "/usr/local/bin/jadx.zip"
     if not os.path.exists(jadx_path):
         print("Instalando o JADX...")
@@ -75,8 +70,6 @@ def install_tools():
         else:
             raise FileNotFoundError("Executável do JADX não encontrado no caminho esperado.")
 
-
-# Função para decompilar o APK
 def decompile_apk(apk_file):
     if apk_file is None:
         return "Nenhum arquivo enviado. Por favor, envie um arquivo APK."
@@ -84,21 +77,18 @@ def decompile_apk(apk_file):
     temp_apk_path = apk_file.name
     output_dir = tempfile.mkdtemp()
     try:
-        # Decompilar usando o Baksmali
         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,
         )
 
-        # Decompilar usando o 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,
         )
 
-        # Coletar arquivos Smali decompilados
         smali_files = {}
         for root, _, files in os.walk(smali_output):
             for file in files:
@@ -106,7 +96,6 @@ def decompile_apk(apk_file):
                     with open(os.path.join(root, file), "r") as f:
                         smali_files[file] = f.read()
 
-        # Coletar arquivos Java decompilados
         java_files = {}
         for root, _, files in os.walk(java_output):
             for file in files:
@@ -114,7 +103,6 @@ def decompile_apk(apk_file):
                     with open(os.path.join(root, file), "r") as f:
                         java_files[file] = f.read()
 
-        # Armazenar resultados no contexto global
         apk_context["smali"] = smali_files
         apk_context["java"] = java_files
 
@@ -123,59 +111,71 @@ def decompile_apk(apk_file):
     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
+        )
+        processed_query = bloom_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return processed_query
 
-# Função para gerar embeddings de textos usando o modelo de indexação
 def get_embeddings(text):
     return indexing_model.encode(text, convert_to_tensor=True)
 
-
-# Função para chat com APK usando CodeBERT
 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."
 
     try:
-        # Preparar os textos
-        smali_texts = [f"{k}\n{v}" for k, v in apk_context["smali"].items()]
-        java_texts = [f"{k}\n{v}" for k, v in apk_context["java"].items()]
-
-        # Obter embedding da query do usuário
-        query_embedding = get_embeddings(user_message)
-
-        # Obter embeddings de Smali e Java
-        smali_embeddings = indexing_model.encode(smali_texts, convert_to_tensor=True)
-        java_embeddings = indexing_model.encode(java_texts, convert_to_tensor=True)
-
-        # Encontrar códigos relevantes usando similaridade
-        smali_scores = util.pytorch_cos_sim(query_embedding, smali_embeddings).squeeze(0)
-        java_scores = util.pytorch_cos_sim(query_embedding, java_embeddings).squeeze(0)
-
-        # Converter para índices inteiros
-        smali_idx = int(torch.argmax(smali_scores).item())
-        java_idx = int(torch.argmax(java_scores).item())
-
-        # Obter os trechos de código
-        smali_result = smali_texts[smali_idx]
-        java_result = java_texts[java_idx]
-
-        # Explicação usando CodeBERT
-        explanation_prompt = f"Explique o seguinte código em detalhes:\n\n{java_result[:500]}"
-        inputs = tokenizer(explanation_prompt, return_tensors="pt", max_length=512, truncation=True, padding=True)
-        with torch.no_grad():
-            outputs = codebert_model.generate(inputs["input_ids"], max_length=512)
-            explanation = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-        response = f"**Código Smali relevante:**\n\n{smali_result[:1000]}\n\n"
-        response += f"**Código Java relevante:**\n\n{java_result[:1000]}\n\n"
-        response += f"**Explicação gerada pelo CodeBERT:**\n\n{explanation}"
+        # 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)
+
+        # 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)
+
+        # 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
+
+        response = ""
+        for idx in top_indices:
+            file_name, relevant_code = combined_texts[idx.item()]
+            
+            # 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)
+            
+            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
 
     except Exception as e:
-        return f"Erro durante a busca e análise: {str(e)}"
-
+        return f"Erro durante a análise: {str(e)}"
 
-# Verificar versão do Java e instalar ferramentas necessárias
+# Configuração e inicialização
 install_tools()
 
 # Interfaces Gradio
@@ -183,17 +183,23 @@ apk_upload_interface = gr.Interface(
     fn=decompile_apk,
     inputs=gr.File(label="Enviar arquivo APK", file_types=[".apk"]),
     outputs="text",
-    title="Analisador de APK",
-    description="Envie um arquivo APK para decompilá-lo em código Smali e Java.",
+    title="Analisador de APK com BLOOM + CodeBERT",
+    description="Envie um arquivo APK para análise avançada com IA."
 )
 
 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=10, label="Resposta do AI"),
-    title="Chat com APK",
-    description="Faça perguntas sobre o código-fonte do APK e receba explicações geradas.",
+    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."
+)
+
+# Interface combinada
+iface = gr.TabbedInterface(
+    [apk_upload_interface, chat_interface],
+    ["Enviar & Analisar", "Análise Inteligente"]
 )
 
-iface = gr.TabbedInterface([apk_upload_interface, chat_interface], ["Enviar & Analisar", "Conversar com AI"])
-iface.launch()
+# Iniciar a interface
+iface.launch()