fix: corrigir OOM e configurações de memória

- Adicionar OMP_NUM_THREADS=1 para evitar erro libgomp
- Configurar PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
- Reduzir batch size de 4 para 2 para economizar memória GPU
- Aumentar gradient_accumulation_steps de 2 para 4
- Adicionar use_cache=False no modelo
- Adicionar configurações de otimização de memória
- Remover código duplicado de TrainingArguments

Dockerfile

@@ -2,6 +2,8 @@ FROM huggingface/transformers-pytorch-gpu:latest
 
 WORKDIR /app
 
+# Instalar dependências do sistema
+# python-is-python3 cria automaticamente o symlink python -> python3
 RUN apt-get update && apt-get install -y --no-install-recommends \
     git \
     python3 \
@@ -9,25 +11,29 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
     python-is-python3 \
     && rm -rf /var/lib/apt/lists/*
 
-RUN python --version && python3 --version
+# Verificar que python está disponível (entrypoint do NVIDIA precisa)
+RUN python --version && \
+    python3 --version && \
+    echo "✅ Python disponível: $(which python)"
 
+# Instalar dependências Python
 COPY requirements.txt .
-RUN python -m pip install --no-cache-dir --upgrade pip && \
-    python -m pip install --no-cache-dir -r requirements.txt
+RUN python3 -m pip install --no-cache-dir --upgrade pip && \
+    python3 -m pip install --no-cache-dir -r requirements.txt
 
+# Copiar scripts de treinamento
 COPY train.py /app/train.py
 COPY app.py /app/app.py
 
-# Diretórios explícitos
-RUN mkdir -p /app/logs /app/outputs
-
-VOLUME ["/app/logs", "/app/outputs"]
-
-EXPOSE 6006
-
+# Configurar variáveis de ambiente padrão (podem ser sobrescritas)
 ENV MODEL_NAME=microsoft/Phi-3-mini-4k-instruct
 ENV DATASET_REPO=beAnalytic/eda-training-dataset
 ENV OUTPUT_REPO=beAnalytic/eda-llm-model
+ENV OMP_NUM_THREADS=1
+ENV PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
 
+# Executar treinamento
+# Usar 'python' (que será o symlink para python3 criado acima)
+# O entrypoint do NVIDIA espera 'python' estar disponível
 CMD ["python", "/app/app.py"]
 

train.py

@@ -15,6 +15,7 @@ from transformers import (
 from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
 from transformers import BitsAndBytesConfig
 from huggingface_hub import login as hf_login, logout as hf_logout
+from huggingface_hub import login as hf_login
 import torch
 import os
 
@@ -68,7 +69,50 @@ else:
 
 # Carregar dataset
 print(f"Carregando dataset: {DATASET_REPO}")
-dataset = load_dataset(DATASET_REPO)
+try:
+    # Tentar carregar o arquivo JSONL expandido diretamente
+    print("Carregando dataset_balanceado_messages_expandido.jsonl...")
+    dataset_raw = load_dataset(
+        DATASET_REPO,
+        data_files={"train": "dataset_balanceado_messages_expandido.jsonl"},
+        verification_mode="no_checks",
+        split="train"
+    )
+    print(f"✅ Arquivo carregado: {len(dataset_raw)} exemplos")
+    
+    # Dividir em train/validation (80/20)
+    print("Dividindo em train/validation (80/20)...")
+    dataset_split = dataset_raw.train_test_split(test_size=0.2, seed=42)
+    dataset = {
+        "train": dataset_split["train"],
+        "validation": dataset_split["test"]
+    }
+    print(f"✅ Dataset preparado: {len(dataset['train'])} exemplos de treino, {len(dataset['validation'])} exemplos de validação")
+    
+except Exception as e:
+    print(f"⚠️ Erro ao carregar arquivo específico: {e}")
+    print("Tentando carregar dataset padrão do repositório...")
+    try:
+        dataset = load_dataset(DATASET_REPO)
+        print(f"✅ Dataset padrão carregado")
+        if "train" not in dataset:
+            # Se não tiver split, criar
+            if len(dataset) == 1:
+                split_name = list(dataset.keys())[0]
+                dataset_raw = dataset[split_name]
+                dataset = dataset_raw.train_test_split(test_size=0.2, seed=42)
+                dataset = {
+                    "train": dataset["train"],
+                    "validation": dataset["test"]
+                }
+                print(f"✅ Dataset dividido: {len(dataset['train'])} treino, {len(dataset['validation'])} validação")
+    except Exception as e2:
+        print(f"❌ Erro ao carregar dataset: {e2}")
+        raise
+
+# Configurar variáveis de ambiente para evitar problemas de memória
+os.environ["OMP_NUM_THREADS"] = "1"
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
 
 # Carregar modelo e tokenizer
 print(f"Carregando modelo: {MODEL_NAME}")
@@ -88,6 +132,7 @@ model = AutoModelForCausalLM.from_pretrained(
     quantization_config=bnb_config,
     device_map="auto",
     trust_remote_code=True,
+    use_cache=False,
 )
 
 # Preparar modelo para LoRA
@@ -107,18 +152,49 @@ model = get_peft_model(model, peft_config)
 
 # Formatar prompts
 def format_prompt(example):
-    system_prompt = (
-        "Você é um analista de dados sênior realizando uma Análise Exploratória de Dados (EDA) "
-        "com rigor estatístico, honestidade analítica e pensamento crítico.\n\n"
-        "Seu objetivo não é gerar insights a qualquer custo, mas avaliar se os dados possuem "
-        "estrutura informativa, comportamento emergente ou apenas relações estruturais triviais."
-    )
-    
-    instruction = example.get("instruction", "")
-    input_text = example.get("input", "")
-    output_text = example.get("output", "")
+    # Suporta dois formatos: messages (ChatML) ou instruction/input/output
+    if "messages" in example:
+        # Formato messages (ChatML)
+        messages = example["messages"]
+        system_content = ""
+        user_content = ""
+        assistant_content = ""
+        
+        for msg in messages:
+            role = msg.get("role", "")
+            content = msg.get("content", "")
+            if role == "system":
+                system_content = content
+            elif role == "user":
+                user_content = content
+            elif role == "assistant":
+                assistant_content = content
+        
+        # Se não houver system, usar padrão
+        if not system_content:
+            system_content = (
+                "Você é um analista de dados sênior realizando uma Análise Exploratória de Dados (EDA) "
+                "com rigor estatístico, honestidade analítica e pensamento crítico.\n\n"
+                "Seu objetivo não é gerar insights a qualquer custo, mas avaliar se os dados possuem "
+                "estrutura informativa, comportamento emergente ou apenas relações estruturais triviais."
+            )
+        
+        prompt = f"<|system|>\n{system_content}\n<|user|>\n{user_content}\n<|assistant|>\n{assistant_content}<|end|>"
+    else:
+        # Formato instruction/input/output (legado)
+        system_prompt = (
+            "Você é um analista de dados sênior realizando uma Análise Exploratória de Dados (EDA) "
+            "com rigor estatístico, honestidade analítica e pensamento crítico.\n\n"
+            "Seu objetivo não é gerar insights a qualquer custo, mas avaliar se os dados possuem "
+            "estrutura informativa, comportamento emergente ou apenas relações estruturais triviais."
+        )
+        
+        instruction = example.get("instruction", "")
+        input_text = example.get("input", "")
+        output_text = example.get("output", "")
+        
+        prompt = f"<|system|>\n{system_prompt}\n<|user|>\n{instruction}\n\n{input_text}\n<|assistant|>\n{output_text}<|end|>"
     
-    prompt = f"<|system|>\n{system_prompt}\n<|user|>\n{instruction}\n\n{input_text}\n<|assistant|>\n{output_text}<|end|>"
     return {"text": prompt}
 
 # Aplicar formatação
@@ -147,32 +223,27 @@ def tokenize_function(examples):
 train_dataset = train_dataset.map(tokenize_function, batched=True, remove_columns=["text"])
 eval_dataset = eval_dataset.map(tokenize_function, batched=True, remove_columns=["text"])
 
-# Criar diretórios manualmente
-os.makedirs("./logs", exist_ok=True)
-os.makedirs("./results", exist_ok=True)
-
 # Configurar argumentos de treinamento
 # push_to_hub_enabled já foi definido acima durante a autenticação
 
 # Argumentos base de treinamento
-# NOTA: logging_dir não é suportado em todas as versões do transformers
-# O TensorBoard funcionará com report_to=["tensorboard"] e salvará logs no output_dir
 training_args_dict = {
     "output_dir": "./results",
     "num_train_epochs": 3,
-    "per_device_train_batch_size": 4,
-    "per_device_eval_batch_size": 4,
+    "per_device_train_batch_size": 2,
+    "per_device_eval_batch_size": 2,
     "learning_rate": 3e-05,
     "warmup_steps": 100,
-    "logging_steps": 1,                   # Para ver métricas rapidamente no TensorBoard
-    "save_steps": 10,                     # Salvar frequentemente para ver progresso
-    "eval_strategy": "steps",             # OBRIGATÓRIO para métricas de validação
-    "eval_steps": 10,                     # Avaliar frequentemente para ver no TensorBoard
+    "logging_steps": 10,
+    "save_steps": 500,
+    "eval_strategy": "steps",
+    "eval_steps": 500,
     "save_total_limit": 3,
     "load_best_model_at_end": True,
     "fp16": True,
-    "gradient_accumulation_steps": 2,
-    "report_to": ["tensorboard"],         # EXPLÍCITO - obrigatório para TensorBoard
+    "gradient_accumulation_steps": 4,
+    "dataloader_pin_memory": False,
+    "ddp_find_unused_parameters": False,
 }
 
 # Adicionar parâmetros do Hub apenas se autenticado