Update train.py

train.py

@@ -1,5 +1,3 @@
-# =================== CÓDIGO FINAL COM SUPORTE PARA XML ===================
-
 import os
 import glob
 import json
@@ -7,62 +5,42 @@ import csv
 import numpy as np
 from tqdm.auto import tqdm
 from sentence_transformers import SentenceTransformer
-import zipfile  # Módulo para lidar com ficheiros .zip
-import xml.etree.ElementTree as ET # Módulo para processar XML
+import zipfile
+import xml.etree.ElementTree as ET
 
-# --- CONFIGURAÇÕES ---
-# Diretório onde o Dockerfile clonou os dados do seu Space
 DATA_DIR = "/app/dados"
-# Diretório para onde os ficheiros serão extraídos
 EXTRACT_DIR = "/app/dados_extraidos"
-# ---------------------
 
 def setup_data():
-    """Procura por ficheiros .zip no diretório de dados e os descompacta."""
-    print("Procurando por ficheiros .zip para descompactar...")
     os.makedirs(EXTRACT_DIR, exist_ok=True)
-
-    zip_files_found = glob.glob(DATA_DIR + "/**/*.zip", recursive=True)
-
-    if not zip_files_found:
-        print("Nenhum ficheiro .zip encontrado. Assumindo que os ficheiros já estão extraídos.")
+    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_filepath in zip_files_found:
-        print(f"Descompactando '{zip_filepath}' para '{EXTRACT_DIR}'...")
-        try:
-            with zipfile.ZipFile(zip_filepath, 'r') as zip_ref:
-                zip_ref.extractall(EXTRACT_DIR)
-            print(f"✅ Ficheiro '{os.path.basename(zip_filepath)}' descompactado com sucesso.")
-        except Exception as e:
-            print(f"⚠️ Falha ao descompactar '{zip_filepath}': {e}")
-
+    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):
-    """Função recursiva para converter um elemento XML num dicionário."""
-    # Se o elemento tem filhos, processa-os
-    if len(element) > 0:
-        d = {}
-        for child in element:
-            child_dict = xml_to_dict(child)
-            # Se a tag já existe, transforma num array
-            if child.tag in d:
-                if not isinstance(d[child.tag], list):
-                    d[child.tag] = [d[child.tag]]
-                d[child.tag].append(child_dict)
-            else:
-                d[child.tag] = child_dict
-        return d
-    # Se não tem filhos, retorna o texto do elemento
-    return element.text
+    d = {}
+    for child in element:
+        child_dict = xml_to_dict(child)
+        if child.tag in d:
+            if not isinstance(d[child.tag], list):
+                d[child.tag] = [d[child.tag]]
+            d[child.tag].append(child_dict)
+        else:
+            d[child.tag] = child_dict
+    if not d:
+        return element.text
+    return d
 
 def serialize_item_to_text(item_dict):
-    """Converte um dicionário (de JSON, CSV ou XML) numa string de texto."""
     parts = []
     if not isinstance(item_dict, dict):
         return str(item_dict)
-
     for key, value in item_dict.items():
         if isinstance(value, dict):
             nested_text = serialize_item_to_text(value)
@@ -75,23 +53,19 @@ def serialize_item_to_text(item_dict):
     return ", ".join(parts)
 
 def main():
-    """Função principal para carregar dados e gerar embeddings."""
     process_dir = setup_data()
-    csv.field_size_limit(100_000_000)
+    csv.field_size_limit(10_000_000)
 
-    # Adiciona .xml à pesquisa de ficheiros
     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)} ficheiros (JSON, CSV, XML) para processar.")
+    print(f"\n🔎 Encontrados {len(all_files)} arquivos (JSON, CSV, XML) para processar.")
 
     if not all_files:
-        print("⚠️ Nenhum ficheiro de dados encontrado. Verifique o conteúdo do seu repositório/zip.")
         return
 
     documents = []
-    for filepath in all_files:
+    for filepath in tqdm(all_files, desc="Processando arquivos"):
         try:
             if filepath.endswith('.json'):
                 with open(filepath, 'r', encoding='utf-8') as f:
@@ -104,16 +78,40 @@ def main():
                 with open(filepath, 'r', encoding='utf-8') as f:
                     reader = csv.DictReader(f)
                     for row in reader: documents.append(serialize_item_to_text(row))
-            # --- LÓGICA PARA PROCESSAR XML ---
             elif filepath.endswith('.xml'):
                 tree = ET.parse(filepath)
                 root = tree.getroot()
                 xml_dict = {root.tag: xml_to_dict(root)}
                 documents.append(serialize_item_to_text(xml_dict))
         except Exception as e:
-            print(f"⚠️ Erro ao processar o ficheiro {filepath}: {e}")
-
-    print(f"\nProcessamento de ficheiros concluído! {len(documents)} documentos foram criados.")
+            print(f"⚠️ Erro ao processar o arquivo {filepath}: {e}")
 
+    print(f"\nProcessamento de arquivos concluído! {len(documents)} documentos foram criados.")
     if not documents:
-        print("Nenhum documento foi lido com sucesso. Encerrando.")
+        return
+
+    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 = 'meus_embeddings_e5_large.npy'
+    
+    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()