Update app.py

app.py

@@ -1,3 +1,5 @@
+!pip install torch==2.0.1 transformers==4.27.1 datasets==2.4.0 wget==3.2 huggingface-hub==0.14.1 beautifulsoup4==4.11.1 requests==2.28.1 matplotlib tqdm python-dotenv diffusers gradio
+
 import os
 import torch
 import torch.nn as nn
@@ -6,12 +8,12 @@ from torch.optim import AdamW
 import matplotlib.pyplot as plt
 import matplotlib.animation as animation
 import time
-import threading
 from tqdm import tqdm
-from transformers import AutoTokenizer, AutoModel, AutoModelForTextToWaveform, TrainingArguments, pipeline
+from transformers import AutoTokenizer, AutoModel, AutoModelForTextToWaveform, 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()
@@ -26,16 +28,15 @@ class UnifiedModel(nn.Module):
         hidden_states = []
         for model in self.models:
             if isinstance(model, nn.Module):
-                outputs = model(inputs)
+                outputs = model(**inputs)
                 hidden_states.append(outputs.last_hidden_state[:, 0, :])
-            elif isinstance(model, DiffusionPipeline) or isinstance(model, pipeline):
-                outputs = model(inputs)
-                hidden_states.append(torch.tensor(outputs))
+            elif isinstance(model, DiffusionPipeline):
+                outputs = model(**inputs)
+                hidden_states.append(torch.tensor(outputs).float())
         concatenated_hidden_states = torch.cat(hidden_states, dim=-1)
         logits = self.classifier(concatenated_hidden_states)
         return logits
 
-
 class SyntheticDataset(Dataset):
     def __init__(self, tokenizers, size=100):
         self.tokenizers = tokenizers
@@ -61,7 +62,6 @@ class SyntheticDataset(Dataset):
     def __getitem__(self, idx):
         return self.data[idx]
 
-
 def push_to_hub(local_dir, repo_name):
     try:
         repo_url = HfApi().create_repo(repo_name, exist_ok=True)
@@ -84,6 +84,55 @@ def push_to_hub(local_dir, repo_name):
     except Exception as e:
         print(f"Error pushing to Hugging Face Hub: {e}")
 
+def load_model(model_name):
+    tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+    model = AutoModel.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:
@@ -107,56 +156,22 @@ def main():
                 "Falconsai/text_summarization",
                 "microsoft/speecht5_tts",
                 "Groq/Llama-3-Groq-70B-Tool-Use",
-                "Groq/Llama-3-Groq-8B-Tool-Use"
-            ]
-
-            # Inicializar los pipelines
-            pipelines_to_unify = [
-                pipeline("text-to-audio", model="facebook/musicgen-melody"),
-                pipeline("text-to-audio", model="facebook/musicgen-large"),
-                pipeline("text-to-audio", model="facebook/musicgen-small"),
-                DiffusionPipeline.from_pretrained("stabilityai/stable-video-diffusion-img2vid-xt-1-1"),
-                pipeline("automatic-speech-recognition", model="openai/whisper-small"),
-                DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev"),
-                DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1"),
-                DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell"),
-                pipeline("text-generation", model="meta-llama/Meta-Llama-3.1-8B"),
-                pipeline("text-generation", model="openbmb/MiniCPM-V-2_6"),
-                pipeline("text-generation", model="bigcode/starcoder"),
-                pipeline("text-to-speech", model="microsoft/speecht5_tts"),
-                pipeline("text-generation", model="WizardLMTeam/WizardCoder-Python-34B-V1.0"),
-                pipeline("text-generation", model="Qwen/Qwen2-72B-Instruct"),
-                pipeline("text-generation", model="google/gemma-2-2b-it"),
-                pipeline("summarization", model="facebook/bart-large-cnn"),
-                pipeline("summarization", model="Falconsai/text_summarization"),
-                DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-dev"),
-                pipeline("text-to-audio", model="facebook/musicgen-small"),
-                pipeline("text-generation", model="Groq/Llama-3-Groq-70B-Tool-Use"),
-                pipeline("text-generation", model="Groq/Llama-3-Groq-8B-Tool-Use")
-            ]
-
-            # Añadir modelos adicionales
-            additional_models = [
+                "Groq/Llama-3-Groq-8B-Tool-Use",
                 "facebook/musicgen-large",
-                "facebook/musicgen-melody"
+                "facebook/musicgen-melody",
+                "black-forest-labs/FLUX.1-schnell",
+                "facebook/musicgen-small",
+                "stabilityai/stable-video-diffusion-img2vid-xt-1-1",
+                "openai/whisper-small",
+                "black-forest-labs/FLUX.1-dev",
+                "stabilityai/stable-diffusion-2-1"
             ]
 
-            # Inicializar los tokenizadores y modelos
+            # Inicializar los modelos y tokenizadores
             tokenizers = {}
             models = []
             for model_name in models_to_train:
-                tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
-
-                if tokenizer.pad_token is None:
-                    tokenizer.add_special_tokens({'pad_token': tokenizer.eos_token})
-
-                model = AutoModel.from_pretrained(model_name)
-                tokenizers[model_name] = tokenizer
-                models.append(model)
-
-            for model_name in additional_models:
-                tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
-                model = AutoModelForTextToWaveform.from_pretrained(model_name)
+                tokenizer, model = load_model(model_name)
                 tokenizers[model_name] = tokenizer
                 models.append(model)
 
@@ -172,7 +187,7 @@ def main():
             train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)
             eval_loader = DataLoader(val_dataset, batch_size=16)
 
-            # Unificar los modelos y pipelines en uno solo
+            # Unificar los modelos en uno solo
             unified_model = UnifiedModel(models)
             unified_model.to(torch.device("cpu"))
 
@@ -199,68 +214,47 @@ def main():
             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()
-
-                    progress_bar.update(1)
-
-                    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")
+            ax.plot(train_losses, label="Training Loss")
+            ax.plot(eval_losses, label="Evaluation Loss")
             ax.legend()
 
             def animate(i):
                 ax.clear()
-                ax.plot(train_losses[:i], label="Train Loss")
-                ax.plot(eval_losses[:i], label="Eval Loss")
+                ax.plot(train_losses, label="Training Loss")
+                ax.plot(eval_losses, label="Evaluation Loss")
+                ax.set_xlabel("Epochs")
+                ax.set_ylabel("Loss")
                 ax.legend()
 
-            ani = animation.FuncAnimation(fig, animate, frames=len(train_losses), blit=False)
+            ani = animation.FuncAnimation(fig, animate, interval=1000)
             plt.show()
 
-            # Subir el modelo unificado a Hugging Face Hub
+            # Guardar el modelo y el tokenizador unificados
+            if not os.path.exists("./outputs/unified_model"):
+                os.makedirs("./outputs/unified_model")
+
+            # Guardar el modelo unificado en un directorio local
             local_dir = "./outputs/unified_model"
+            torch.save(unified_model.state_dict(), os.path.join(local_dir, "pytorch_model.bin"))
+
+            # Guardar el tokenizador en un directorio local
+            for name, tokenizer in tokenizers.items():
+                tokenizer.save_pretrained(local_dir)
+
+            # 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}")