Update app.py

app.py

@@ -7,10 +7,11 @@ 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
+import gradio as gr
 
 # Cargar variables de entorno
 load_dotenv()
@@ -83,9 +84,54 @@ 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):
+    model.train()
+    epoch = 0
+    total_steps = len(train_loader)
+    for step, batch in enumerate(train_loader):
+        start_time = time.time()
+        input_ids = [batch[f"input_ids_{name}"].to("cpu") for name in tokenizers.keys()]
+        attention_mask = [batch[f"attention_mask_{name}"].to("cpu") for name in tokenizers.keys()]
+        labels = batch["label"].to("cpu")
+        optimizer.zero_grad()
+        outputs = model(input_ids)
+        loss = nn.CrossEntropyLoss()(outputs, labels)
+        loss.backward()
+        optimizer.step()
+
+        elapsed_time = time.time() - start_time
+        estimated_total_time = total_steps * (elapsed_time / (step + 1))
+        estimated_remaining_time = estimated_total_time - elapsed_time
+
+        if step % args.logging_steps == 0:
+            train_losses.append(loss.item())
+            print(f"Step {step}/{total_steps}, Loss: {loss.item()}, Estimated remaining time: {estimated_remaining_time:.2f} seconds")
+
+    epoch += 1
+    model.eval()
+    eval_loss = 0
+    with torch.no_grad():
+        for batch in eval_loader:
+            input_ids = [batch[f"input_ids_{name}"].to("cpu") for name in tokenizers.keys()]
+            attention_mask = [batch[f"attention_mask_{name}"].to("cpu") for name in tokenizers.keys()]
+            labels = batch["label"].to("cpu")
+            outputs = model(input_ids)
+            loss = nn.CrossEntropyLoss()(outputs, labels)
+            eval_loss += loss.item()
+
+    eval_loss /= len(eval_loader)
+    eval_losses.append(eval_loss)
+    print(f"Epoch {epoch}/{args.num_train_epochs}, Evaluation Loss: {eval_loss}")
+
+def gradio_interface(input_text):
+    # Define the Gradio interface function
+    tokenized_inputs = {name: tokenizer.encode(input_text, return_tensors="pt") for name, tokenizer in tokenizers.items()}
+    model_output = unified_model(tokenized_inputs)
+    return model_output
+
 def main():
     while True:
         try:
@@ -149,65 +195,25 @@ 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 = []
 
-            def train(model, train_loader, eval_loader, args):
-                model.train()
-                epoch = 0
-                total_steps = len(train_loader)
-                for step, batch in enumerate(train_loader):
-                    start_time = time.time()
-                    input_ids = [batch[f"input_ids_{name}"].to("cpu") for name in tokenizers.keys()]
-                    attention_mask = [batch[f"attention_mask_{name}"].to("cpu") for name in tokenizers.keys()]
-                    labels = batch["label"].to("cpu")
-                    optimizer.zero_grad()
-                    outputs = model(input_ids)
-                    loss = nn.CrossEntropyLoss()(outputs, labels)
-                    loss.backward()
-                    optimizer.step()
-
-                    elapsed_time = time.time() - start_time
-                    estimated_total_time = total_steps * (elapsed_time / (step + 1))
-                    estimated_remaining_time = estimated_total_time - elapsed_time
-
-                    if step % args.logging_steps == 0:
-                        train_losses.append(loss.item())
-                        print(f"Step {step}/{total_steps}, Loss: {loss.item()}, Estimated remaining time: {estimated_remaining_time:.2f} seconds")
-
-                epoch += 1
-                model.eval()
-                eval_loss = 0
-                with torch.no_grad():
-                    for batch in eval_loader:
-                        input_ids = [batch[f"input_ids_{name}"].to("cpu") for name in tokenizers.keys()]
-                        attention_mask = [batch[f"attention_mask_{name}"].to("cpu") for name in tokenizers.keys()]
-                        labels = batch["label"].to("cpu")
-                        outputs = model(input_ids)
-                        loss = nn.CrossEntropyLoss()(outputs, labels)
-                        eval_loss += loss.item()
-
-                eval_loss /= len(eval_loader)
-                eval_losses.append(eval_loss)
-                print(f"Epoch {epoch}/{args.num_train_epochs}, Evaluation Loss: {eval_loss}")
-
+            # Entrenar el modelo
             train(unified_model, train_loader, eval_loader, training_args)
 
-            # Visualizar pérdidas durante el entrenamiento
+            # Visualizar pérdidas
             fig, ax = plt.subplots()
             ax.set_xlabel("Epochs")
             ax.set_ylabel("Loss")
@@ -241,6 +247,10 @@ def main():
             # Subir el modelo y el tokenizador a Hugging Face
             push_to_hub(local_dir, repo_name="Ffftdtd5dtft/my_model")
 
+            # Configurar y lanzar la interfaz Gradio
+            interface = gr.Interface(fn=gradio_interface, inputs="text", outputs="text")
+            interface.launch()
+
             break
         except Exception as e:
             print(f"Error: {e}")