update

train.py

@@ -1,4 +1,4 @@
-# =================== CÓDIGO COMPLETO E FINAL ===================
+# =================== CÓDIGO FINAL COM LÓGICA DE DESCOMPACTAÇÃO ===================
 
 import os
 import glob
@@ -7,12 +7,39 @@ import csv
 import numpy as np
 from tqdm.auto import tqdm
 from sentence_transformers import SentenceTransformer
+import zipfile # Módulo para lidar com arquivos .zip
 
 # --- CONFIGURAÇÕES ---
-# O diretório onde o Dockerfile clonou os dados do próprio Space
+# Diretório onde o Dockerfile clonou os dados do seu Space
 DATA_DIR = "/app/dados"
+# Diretório para onde os arquivos serão extraídos
+EXTRACT_DIR = "/app/dados_extraidos"
 # ---------------------
 
+def setup_data():
+    """Procura por arquivos .zip no diretório de dados e os descompacta."""
+    print("Procurando por arquivos .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 arquivo .zip encontrado. Assumindo que os arquivos já estão extraídos.")
+        # Se não houver zips, usamos o diretório de dados original
+        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"✅ Arquivo '{os.path.basename(zip_filepath)}' descompactado com sucesso.")
+        except Exception as e:
+            print(f"⚠️ Falha ao descompactar '{zip_filepath}': {e}")
+            
+    # Retorna o diretório onde os arquivos foram extraídos
+    return EXTRACT_DIR
+
 def serialize_item_to_text(item_dict):
     """Converte um dicionário em uma string de texto."""
     parts = []
@@ -29,16 +56,19 @@ def serialize_item_to_text(item_dict):
 
 def main():
     """Função principal para carregar dados e gerar embeddings."""
+    # Descompacta os dados e obtém o diretório com os arquivos processáveis
+    process_dir = 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)
+    # Procura recursivamente por todos os arquivos .json e .csv no diretório de extração
+    all_files = glob.glob(process_dir + "/**/*.json", recursive=True) + \
+                glob.glob(process_dir + "/**/*.csv", recursive=True)
 
-    print(f"🔎 Encontrados {len(all_files)} arquivos para processar no repositório.")
+    print(f"\n🔎 Encontrados {len(all_files)} arquivos para processar após descompactação.")
     
     if not all_files:
-        print("⚠️ Nenhum arquivo .csv ou .json encontrado. Verifique se os dados estão no repositório.")
+        print("⚠️ Nenhum arquivo .csv ou .json encontrado. Verifique o conteúdo do seu .zip.")
         return
         
     documents = []
@@ -68,7 +98,6 @@ def main():
     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,
@@ -93,4 +122,4 @@ def main():
 if __name__ == "__main__":
     main()
 
-# =================================================================
+# =================================================================================