Update app.py

app.py

@@ -1,3 +1,4 @@
+!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
 
 import os
 import torch
@@ -7,9 +8,8 @@ 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, 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
@@ -27,16 +27,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
@@ -62,7 +61,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)
@@ -85,6 +83,10 @@ 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 main():
     while True:
@@ -108,49 +110,25 @@ 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")
+                "Groq/Llama-3-Groq-8B-Tool-Use",
+                "facebook/musicgen-large",
+                "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 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)
+                tokenizer, model = load_model(model_name)
                 tokenizers[model_name] = tokenizer
                 models.append(model)
 
-            # Agregar pipelines como modelos
-            models.extend(pipelines_to_unify)
-
             # Crear un dataset sintético para entrenamiento y evaluación
             synthetic_dataset = SyntheticDataset(tokenizers, size=100)
 
@@ -163,7 +141,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"))
 
@@ -193,45 +171,41 @@ def main():
             def train(model, train_loader, eval_loader, args):
                 model.train()
                 epoch = 0
-                total_steps = args.num_train_epochs * len(train_loader)
-                progress_bar = tqdm(total=total_steps, desc="Training")
-
-                while epoch < args.num_train_epochs:
+                total_steps = len(train_loader)
+                for step, batch in enumerate(train_loader):
                     start_time = time.time()
-                    for step, batch in enumerate(train_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")
+                    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")
-                        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}")
+                        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}")
 
             train(unified_model, train_loader, eval_loader, training_args)
 
@@ -239,19 +213,34 @@ def main():
             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")
 
             break
@@ -260,4 +249,4 @@ def main():
             time.sleep(2)
 
 if __name__ == "__main__":
-    main()
+    main()