Update app.py

app.py

@@ -7,7 +7,7 @@ import matplotlib.pyplot as plt
 import matplotlib.animation as animation
 import time
 from tqdm import tqdm
-from transformers import AutoTokenizer, AutoModel, AutoModelForTextToWaveform, TrainingArguments
+from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
 from diffusers import DiffusionPipeline
 from huggingface_hub import login, HfApi, Repository
 from dotenv import load_dotenv
@@ -84,7 +84,7 @@ def push_to_hub(local_dir, repo_name):
 
 def load_model(model_name):
     tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
-    model = AutoModel.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
     return tokenizer, model
 
 def train(model, train_loader, eval_loader, args):
@@ -195,20 +195,18 @@ def main():
 
             # Definir los argumentos de entrenamiento
             training_args = TrainingArguments(
-                output_dir="outputs/unified_model",
-                evaluation_strategy="epoch",
-                learning_rate=9e-4,
                 per_device_train_batch_size=2,
                 per_device_eval_batch_size=16,
-                num_train_epochs=1,  # Reduced epochs for quick training
-                weight_decay=0.01,
-                logging_steps=10,  # More frequent logging for quicker feedback
-                optim="adamw_hf"
+                num_train_epochs=1,
+                logging_steps=10,
+                save_steps=10,
+                evaluation_strategy="steps"
             )
 
             # Definir el optimizador
-            optimizer = AdamW(unified_model.parameters(), lr=training_args.learning_rate)
+            optimizer = AdamW(unified_model.parameters(), lr=5e-5)
 
+            # Listas para almacenar las pérdidas
             train_losses = []
             eval_losses = []