update

train.py

@@ -1,7 +1,6 @@
-# =================== CÓDIGO COMPLETO E ATUALIZADO ===================
+# =================== CÓDIGO COMPLETO E FINAL ===================
 
 import os
-import zipfile
 import glob
 import json
 import csv
@@ -10,24 +9,10 @@ from tqdm.auto import tqdm
 from sentence_transformers import SentenceTransformer
 
 # --- CONFIGURAÇÕES ---
-# Lembre-se de colocar o nome correto do seu arquivo .zip!
-ZIP_FILENAME = "seus-dados.zip"
-EXTRACT_DIR = "/app/dados_extraidos"
+# O diretório onde o Dockerfile clonou os dados do próprio Space
+DATA_DIR = "/app/dados"
 # ---------------------
 
-def setup_data():
-    """Descompacta os dados se o diretório não existir."""
-    if not os.path.exists(EXTRACT_DIR) and os.path.exists(ZIP_FILENAME):
-        print(f"Descompactando '{ZIP_FILENAME}'...")
-        os.makedirs(EXTRACT_DIR, exist_ok=True)
-        with zipfile.ZipFile(ZIP_FILENAME, 'r') as zip_ref:
-            zip_ref.extractall(EXTRACT_DIR)
-        print("✅ Dados descompactados.")
-    elif not os.path.exists(ZIP_FILENAME):
-        print(f"⚠️ Arquivo '{ZIP_FILENAME}' não encontrado. Pulando descompactação.")
-    else:
-        print("✅ Dados já parecem estar descompactados.")
-
 def serialize_item_to_text(item_dict):
     """Converte um dicionário em uma string de texto."""
     parts = []
@@ -44,8 +29,68 @@ def serialize_item_to_text(item_dict):
 
 def main():
     """Função principal para carregar dados e gerar embeddings."""
-    setup_data()
     csv.field_size_limit(10_000_000)
 
+    # Procura recursivamente por todos os arquivos .json e .csv no diretório de dados
+    all_files = glob.glob(DATA_DIR + "/**/*.json", recursive=True) + \
+                glob.glob(DATA_DIR + "/**/*.csv", recursive=True)
+
+    print(f"🔎 Encontrados {len(all_files)} arquivos para processar no repositório.")
+    
+    if not all_files:
+        print("⚠️ Nenhum arquivo .csv ou .json encontrado. Verifique se os dados estão no repositório.")
+        return
+        
     documents = []
-    all_files = glob.glob(EXTRACT_DIR + "/**/*.
+    for filepath in all_files:
+        try:
+            if filepath.endswith('.json'):
+                with open(filepath, 'r', encoding='utf-8') as f:
+                    data = json.load(f)
+                    if isinstance(data, list):
+                        for item in data: documents.append(serialize_item_to_text(item))
+                    else:
+                        documents.append(serialize_item_to_text(data))
+            elif filepath.endswith('.csv'):
+                with open(filepath, 'r', encoding='utf-8') as f:
+                    reader = csv.DictReader(f)
+                    for row in reader: documents.append(serialize_item_to_text(row))
+        except Exception as e:
+            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
+
+    # Define o caminho do cache e carrega o modelo
+    cache_path = os.environ.get('SENTENCE_TRANSFORMERS_HOME', '/app/cache/torch')
+    
+    print("Carregando modelo avançado: intfloat/e5-mistral-7b-instruct")
+    print("Isso pode levar vários minutos, pois o modelo é grande.")
+    model = SentenceTransformer(
+        'intfloat/e5-mistral-7b-instruct',
+        cache_folder=cache_path,
+        trust_remote_code=True
+    )
+    
+    batch_size = 64
+    output_filename = 'meus_embeddings_finais.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)):
+        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()
+
+# =================================================================