Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import platform
 from dotenv import load_dotenv
 import torch
-from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments
+from transformers import GPT2LMHeadModel, GPT2Tokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling
 from datasets import load_dataset, concatenate_datasets
 from huggingface_hub import login
 import time
@@ -44,7 +44,7 @@ async def root():
 def load_and_train():
     model_name = 'gpt2'
     tokenizer = GPT2Tokenizer.from_pretrained(model_name)
-    model = GPT2LMHeadModel.from_pretrained(model_name)
+    model = GPT2LMHeadModel.from_pretrained(model_name, return_dict=True)
 
     # Asignar el pad_token al eos_token
     tokenizer.pad_token = tokenizer.eos_token
@@ -128,13 +128,14 @@ def load_and_train():
 
     # Función de tokenización basada en el campo 'text'
     def tokenize_function(examples):
-        return tokenizer(
+        tokenized = tokenizer(
             examples['text'], 
             truncation=True, 
             padding='max_length', 
             max_length=512
-            # clean_up_tokenization_spaces=True  # Eliminado porque no es reconocido
         )
+        tokenized['labels'] = tokenized['input_ids'].copy()
+        return tokenized
 
     # Tokenizar el dataset
     tokenized_dataset = combined_dataset.map(
@@ -142,6 +143,12 @@ def load_and_train():
         batched=True
     )
 
+    # Configurar el Data Collator
+    data_collator = DataCollatorForLanguageModeling(
+        tokenizer=tokenizer,
+        mlm=False  # Para modelado de lenguaje causal
+    )
+
     # Configurar argumentos de entrenamiento
     training_args = TrainingArguments(
         output_dir=os.path.join(cache_dir, 'results'),  # Almacenar temporalmente en RAM
@@ -164,6 +171,7 @@ def load_and_train():
         model=model,
         args=training_args,
         train_dataset=tokenized_dataset,
+        data_collator=data_collator,
     )
 
     while True: