Todas as partidas

app.py → train.py

@@ -1,42 +1,29 @@
-# --- 1. IMPORTS ---
-# Imports do seu train.py e do Gradio
 import os
 import glob
 import json
 import csv
 import numpy as np
+from tqdm.auto import tqdm
 from sentence_transformers import SentenceTransformer
 import zipfile
 import xml.etree.ElementTree as ET
-import gradio as gr
-import shutil
 
-# --- 2. CONFIGURAÇÕES E CONSTANTES ---
-# Caminhos relativos são melhores para portabilidade no Hugging Face Spaces
-DATA_DIR = "dados"
-EXTRACT_DIR = os.path.join(DATA_DIR, "dados_extraidos")
-
-# ATIVA O ARMAZENAMENTO PERSISTENTE NAS CONFIGURAÇÕES DO SEU SPACE!
-# Se ativado, mude o caminho para algo como "/data/meus_embeddings.npy"
-OUTPUT_FILENAME = "meus_embeddings_e5_large.npy" 
-
-# --- 3. SUAS FUNÇÕES DE PROCESSAMENTO (DO TRAIN.PY) ---
-# Copiamos suas funções de ajuda diretamente para cá.
+DATA_DIR = "/app/dados"
+EXTRACT_DIR = "/app/dados_extraidos"
 
 def setup_data():
-    """Descompacta os arquivos .zip e retorna o diretório de processamento."""
     os.makedirs(EXTRACT_DIR, exist_ok=True)
-    zip_files = glob.glob(os.path.join(DATA_DIR, "*.zip")) # Simplificado para buscar zips na pasta 'dados'
+    zip_files = glob.glob(DATA_DIR + "/**/*.zip", recursive=True)
     if not zip_files:
         print("Nenhum arquivo .zip encontrado, usando o diretório de dados principal.")
         return DATA_DIR
     for zip_path in zip_files:
+        print(f"Descompactando {zip_path}...")
         with zipfile.ZipFile(zip_path, 'r') as zf:
             zf.extractall(EXTRACT_DIR)
     return EXTRACT_DIR
 
 def xml_to_dict(element):
-    """Converte um elemento XML para um dicionário Python."""
     d = {}
     for child in element:
         child_dict = xml_to_dict(child)
@@ -51,7 +38,6 @@ def xml_to_dict(element):
     return d
 
 def serialize_item_to_text(item_dict):
-    """Converte um dicionário (de JSON, CSV, etc.) para uma string de texto plano."""
     parts = []
     if not isinstance(item_dict, dict):
         return str(item_dict)
@@ -66,30 +52,21 @@ def serialize_item_to_text(item_dict):
             parts.append(f"{key}: {value}")
     return ", ".join(parts)
 
-# --- 4. FUNÇÃO ORQUESTRADORA (LÓGICA PRINCIPAL) ---
-# Esta função substitui a sua função `main()` e é chamada pelo Gradio.
-# Ela usa `yield` para enviar atualizações de progresso para a interface.
-
-def run_full_process():
-    """Executa o pipeline completo e envia o progresso para a UI."""
-    
-    # --- ETAPA 1: SETUP E PROCESSAMENTO DE ARQUIVOS ---
-    yield "Iniciando... Descompactando arquivos..."
+def main():
     process_dir = setup_data()
-    
     csv.field_size_limit(10_000_000)
-    all_files = glob.glob(os.path.join(process_dir, "**/*.json"), recursive=True) + \
-                glob.glob(os.path.join(process_dir, "**/*.csv"), recursive=True) + \
-                glob.glob(os.path.join(process_dir, "**/*.xml"), recursive=True)
-    
-    yield f"🔎 Encontrados {len(all_files)} arquivos para processar."
+
+    all_files = glob.glob(process_dir + "/**/*.json", recursive=True) + \
+                glob.glob(process_dir + "/**/*.csv", recursive=True) + \
+                glob.glob(process_dir + "/**/*.xml", recursive=True)
+    print(f"\n🔎 Encontrados {len(all_files)} arquivos (JSON, CSV, XML) para processar.")
+
+    if not all_files:
+        return
 
     documents = []
-    for idx, filepath in enumerate(all_files):
+    for filepath in tqdm(all_files, desc="Processando arquivos"):
         try:
-            # Mostra o progresso na interface em vez de usar tqdm
-            yield f"Processando arquivo {idx + 1}/{len(all_files)}: {os.path.basename(filepath)}"
-            
             if filepath.endswith('.json'):
                 with open(filepath, 'r', encoding='utf-8') as f:
                     data = json.load(f)
@@ -107,17 +84,34 @@ def run_full_process():
                 xml_dict = {root.tag: xml_to_dict(root)}
                 documents.append(serialize_item_to_text(xml_dict))
         except Exception as e:
-            yield f"⚠️ Erro ao processar {os.path.basename(filepath)}: {e}"
+            print(f"⚠️ Erro ao processar o arquivo {filepath}: {e}")
 
-    yield f"Processamento de arquivos concluído! {len(documents)} documentos criados."
+    print(f"\nProcessamento de arquivos concluído! {len(documents)} documentos foram criados.")
     if not documents:
-        yield "Nenhum documento encontrado para gerar embeddings. Processo encerrado."
         return
 
-    # --- ETAPA 2: GERAÇÃO DE EMBEDDINGS ---
-    yield "Carregando modelo de alta performance: intfloat/multilingual-e5-large..."
+    cache_path = os.environ.get('SENTENCE_TRANSFORMERS_HOME', '/app/cache/torch')
+    
+    print("Carregando modelo de alta performance: intfloat/multilingual-e5-large")
+    model = SentenceTransformer(
+        'intfloat/multilingual-e5-large',
+        cache_folder=cache_path
+    )
+    
+    batch_size = 128
+    output_filename = '/app/output/meus_embeddings_e5_large.npy'
     
-    # Use um cache dentro do seu Space para não baixar o modelo toda vez
-    cache_path = './model_cache'
-    os.makedirs(cache_path, exist_ok=True)
-    model = SentenceTransformer('intfloat/multilingual-
+    if os.path.exists(output_filename):
+        os.remove(output_filename)
+
+    print(f"🚀 Iniciando geração de embeddings (lotes de {batch_size}).")
+    for i in tqdm(range(0, len(documents), batch_size), desc="Gerando Embeddings"):
+        batch = documents[i:i + batch_size]
+        batch_embeddings = model.encode(batch, show_progress_bar=False)
+        with open(output_filename, 'ab') as f_out:
+            np.save(f_out, batch_embeddings)
+
+    print(f"✅ Processo finalizado! Embeddings salvos em '{output_filename}'.")
+
+if __name__ == "__main__":
+        main()