Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import platform
 from dotenv import load_dotenv
 import torch
 from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments
@@ -8,12 +9,23 @@ import time
 import uvicorn
 from fastapi import FastAPI
 import threading
+import logging
+
+# Configurar logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler("training.log"),
+        logging.StreamHandler()
+    ]
+)
 
 # Cargar las variables de entorno
 load_dotenv()
 huggingface_token = os.getenv('HUGGINGFACE_TOKEN')
 if huggingface_token is None:
-    raise ValueError("HUGGINGFACE_TOKEN not found in environment variables.")
+    raise ValueError("HUGGINGFACE_TOKEN no encontrado en las variables de entorno.")
 
 # Iniciar sesión en Hugging Face
 login(token=huggingface_token)
@@ -30,43 +42,72 @@ def load_and_train():
     tokenizer = GPT2Tokenizer.from_pretrained(model_name)
     model = GPT2LMHeadModel.from_pretrained(model_name)
 
+    # Asignar el pad_token al eos_token
+    tokenizer.pad_token = tokenizer.eos_token
+
+    # Redimensionar las embeddings del modelo para incluir el pad_token
+    model.resize_token_embeddings(len(tokenizer))
+
+    # Verificar dispositivo
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+    logging.info(f"Entrenando en: {device}")
+
+    # Determinar cache_dir
+    if platform.system() == "Linux":
+        cache_dir = '/dev/shm'
+    else:
+        cache_dir = './cache'
+
+    # Crear el directorio de caché si no existe
+    os.makedirs(cache_dir, exist_ok=True)
+
     # Intentar cargar los datasets con manejo de errores
     try:
-        dataset_humanizado = load_dataset('daily_dialog', split='train', cache_dir='/dev/shm', trust_remote_code=True)
-        dataset_codigo = load_dataset('code_search_net', split='train', cache_dir='/dev/shm', trust_remote_code=True)
+        dataset_humanizado = load_dataset('daily_dialog', split='train', cache_dir=cache_dir, trust_remote_code=True)
+        dataset_codigo = load_dataset('code_search_net', split='train', cache_dir=cache_dir, trust_remote_code=True)
     except Exception as e:
-        print(f"Error al cargar los datasets: {e}")
-        # Si hay un error, podrías intentar cargar un dataset alternativo o reintentar después de un tiempo
+        logging.error(f"Error al cargar los datasets: {e}")
+        # Intentar cargar un dataset alternativo
         time.sleep(60)  # Esperar 60 segundos antes de reintentar
         try:
-            dataset_humanizado = load_dataset('alternative_dataset', split='train', cache_dir='/dev/shm', trust_remote_code=True)
+            dataset_humanizado = load_dataset('alternative_dataset', split='train', cache_dir=cache_dir, trust_remote_code=True)
+            dataset_codigo = load_dataset('alternative_code_dataset', split='train', cache_dir=cache_dir, trust_remote_code=True)
         except Exception as e:
-            print(f"Error al cargar el dataset alternativo: {e}")
+            logging.error(f"Error al cargar el dataset alternativo: {e}")
             return
 
-    print("Daily Dialog columns:", dataset_humanizado.column_names)
-    print("Code Search Net columns:", dataset_codigo.column_names)
+    logging.info("Daily Dialog columnas: %s", dataset_humanizado.column_names)
+    logging.info("Code Search Net columnas: %s", dataset_codigo.column_names)
 
     # Combinar los datasets en memoria
     combined_dataset = concatenate_datasets([dataset_humanizado, dataset_codigo])
 
-    print("Combined dataset columns:", combined_dataset.column_names)
+    logging.info("Dataset combinado columnas: %s", combined_dataset.column_names)
 
     # Función de tokenización en RAM
     def tokenize_function(examples):
+        text = ""
         if 'dialog' in examples:
-            return tokenizer(examples['dialog'], truncation=True, padding='max_length', max_length=512)
+            text = examples['dialog']
         elif 'docstring' in examples:
-            return tokenizer(examples['docstring'], truncation=True, padding='max_length', max_length=512)
+            text = examples['docstring']
         elif 'code' in examples:
-            return tokenizer(examples['code'], truncation=True, padding='max_length', max_length=512)
+            text = examples['code']
+        if text:
+            return tokenizer(text, truncation=True, padding='max_length', max_length=512)
         return {}
 
     # Tokenizar y mantener todo en RAM
-    tokenized_dataset = combined_dataset.map(tokenize_function, batched=True, cache_file_name='/dev/shm/tokenized_dataset.arrow')
+    tokenized_dataset = combined_dataset.map(
+        tokenize_function, 
+        batched=True, 
+        cache_dir=cache_dir
+    )
 
+    # Configurar argumentos de entrenamiento
     training_args = TrainingArguments(
-        output_dir='/dev/shm/results',  # Almacenar temporalmente en RAM
+        output_dir=os.path.join(cache_dir, 'results'),  # Almacenar temporalmente en RAM
         per_device_train_batch_size=4,
         per_device_eval_batch_size=4,
         num_train_epochs=1,
@@ -78,7 +119,6 @@ def load_and_train():
         warmup_ratio=0.1,
         evaluation_strategy="epoch",
         lr_scheduler_type="linear",
-        save_steps=500,  # Guardar menos frecuentemente para evitar escritura
         save_strategy="epoch",  # Guardar solo al final de cada epoch
     )
 
@@ -92,14 +132,15 @@ def load_and_train():
         try:
             trainer.train()
             # Subir el modelo a Hugging Face desde la RAM
-            model.push_to_hub('Yhhxhfh/nombre_de_tu_modelo', repo_type='model', commit_message="Actualización del modelo")
-            tokenizer.push_to_hub('Yhhxhfh/nombre_de_tu_modelo', repo_type='model', commit_message="Actualización del tokenizador")
-            time.sleep(300)
+            model.push_to_hub('Yhhxhfh/nombre_de_tu_modelo', commit_message="Actualización del modelo")
+            tokenizer.push_to_hub('Yhhxhfh/nombre_de_tu_modelo', commit_message="Actualización del tokenizador")
+            logging.info("Modelo y tokenizador subidos exitosamente.")
+            time.sleep(300)  # Esperar 5 minutos antes de la siguiente iteración
         except Exception as e:
-            print(f"Error durante el entrenamiento: {e}. Reiniciando el proceso de entrenamiento...")
+            logging.error(f"Error durante el entrenamiento: {e}. Reiniciando el proceso de entrenamiento...")
             time.sleep(10)
 
 if __name__ == "__main__":
     # Correr FastAPI en un hilo separado
-    threading.Thread(target=lambda: uvicorn.run(app, host="0.0.0.0", port=7860)).start()
+    threading.Thread(target=lambda: uvicorn.run(app, host="0.0.0.0", port=7860), daemon=True).start()
     load_and_train()