Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import shutil
 import os
 from zipfile import ZipFile
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 # Diretório temporário para os uploads
@@ -11,9 +11,9 @@ UPLOAD_DIR = "uploaded_files"
 # Variável global para armazenar o código dos arquivos .java
 java_files_content = {}
 
-# Carregar o modelo e o tokenizador do CodeBERT localmente
+# Carregar o modelo e o tokenizador do CodeBERT
 tokenizer = AutoTokenizer.from_pretrained("microsoft/codebert-base")
-model = AutoModelForSequenceClassification.from_pretrained("microsoft/codebert-base")
+model = AutoModelForCausalLM.from_pretrained("microsoft/codebert-base")
 
 # Função para processar a pasta enviada
 def process_uploaded_folder(zip_path):
@@ -54,6 +54,7 @@ def search_in_code(search_term):
     results = []
     for filename, content in java_files_content.items():
         if search_term.lower() in content.lower():
+            # Adiciona o nome do arquivo e um trecho do código onde foi encontrado
             results.append(f"Arquivo: {filename}\n{content[:300]}...")  # Exibe os primeiros 300 caracteres do código
     
     if not results:
@@ -66,12 +67,11 @@ def generate_codebert_response(prompt):
     # Tokenize the input
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, padding=True, max_length=512)
     with torch.no_grad():
-        outputs = model(**inputs)
-    
-    # Decode the output (for simplicity, assuming a classification task, but could be adjusted)
-    response = outputs.logits.argmax(dim=-1).item()  # Simplified for classification, adjust as needed
-    
-    return f"Resposta do modelo: {response}"
+        outputs = model.generate(**inputs, max_length=512, num_return_sequences=1)
+
+    # Decodificar a resposta gerada
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
 
 # Interface de upload
 def upload_and_analyze(zip_file_path):