Update app.py

app.py

@@ -6,15 +6,12 @@ import torch
 from transformers import AutoTokenizer, DataCollatorWithPadding, AutoModelForSequenceClassification, AdamW
 from datasets import load_dataset
 from evaluate import load as load_metric
-from torch.utils.data import DataLoader
+from torch.utils.data import DataLoader, random_split
 import random
 
 DEVICE = torch.device("cpu")
-NUM_ROUNDS = 3
 
-
-
-def load_data(dataset_name):
+def load_data(dataset_name, train_size=20, test_size=20):
     raw_datasets = load_dataset(dataset_name)
     raw_datasets = raw_datasets.shuffle(seed=42)
     del raw_datasets["unsupervised"]
@@ -24,19 +21,16 @@ def load_data(dataset_name):
     def tokenize_function(examples):
         return tokenizer(examples["text"], truncation=True)
 
-    train_population = random.sample(range(len(raw_datasets["train"])), 20)
-    test_population = random.sample(range(len(raw_datasets["test"])), 20)
-
     tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
-    tokenized_datasets["train"] = tokenized_datasets["train"].select(train_population)
-    tokenized_datasets["test"] = tokenized_datasets["test"].select(test_population)
-
     tokenized_datasets = tokenized_datasets.remove_columns("text")
     tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
 
+    train_dataset = tokenized_datasets["train"].select(random.sample(range(len(tokenized_datasets["train"])), train_size))
+    test_dataset = tokenized_datasets["test"].select(random.sample(range(len(tokenized_datasets["test"])), test_size))
+
     data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
-    trainloader = DataLoader(tokenized_datasets["train"], shuffle=True, batch_size=32, collate_fn=data_collator)
-    testloader = DataLoader(tokenized_datasets["test"], batch_size=32, collate_fn=data_collator)
+    trainloader = DataLoader(train_dataset, shuffle=True, batch_size=32, collate_fn=data_collator)
+    testloader = DataLoader(test_dataset, batch_size=32, collate_fn=data_collator)
 
     return trainloader, testloader
 
@@ -54,8 +48,8 @@ def train(net, trainloader, epochs):
 
 def test(net, testloader):
     metric = load_metric("accuracy")
-    loss = 0
     net.eval()
+    loss = 0
     for batch in testloader:
         batch = {k: v.to(DEVICE) for k, v in batch.items()}
         with torch.no_grad():
@@ -64,29 +58,16 @@ def test(net, testloader):
         loss += outputs.loss.item()
         predictions = torch.argmax(logits, dim=-1)
         metric.add_batch(predictions=predictions, references=batch["labels"])
-    loss /= len(testloader.dataset)
+    loss /= len(testloader)
     accuracy = metric.compute()["accuracy"]
     return loss, accuracy
 
-
-
-
-
-from transformers import Wav2Vec2Processor, HubertForSequenceClassification
-import torch
-
-    
 def main():
-    st.write("## Federated Learning with dynamic models and datasets for mobile devices")
-    dataset_name = st.selectbox("Dataset", ["imdb","audio_instruction_task", "amazon_polarity", "ag_news"])
-    model_name = st.selectbox("Model", ["bert-base-uncased","facebook/hubert-base-ls960", "distilbert-base-uncased"])
+    st.write("## Federated Learning with Dynamic Models and Datasets for Mobile Devices")
+    dataset_name = st.selectbox("Dataset", ["imdb", "amazon_polarity", "ag_news"])
+    model_name = st.selectbox("Model", ["bert-base-uncased", "distilbert-base-uncased"])
 
     net = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2).to(DEVICE)
-    # processor = Wav2Vec2Processor.from_pretrained(model_name)
-    # net = HubertForSequenceClassification.from_pretrained(model_name, num_labels=2).to(DEVICE)
-
-    # feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name)
-    # net = HubertForSequenceClassification.from_pretrained(model_name, num_labels=2).to(DEVICE)
 
     NUM_CLIENTS = st.slider("Number of Clients", min_value=1, max_value=10, value=2)
     NUM_ROUNDS = st.slider("Number of Rounds", min_value=1, max_value=10, value=3)
@@ -95,28 +76,42 @@ def main():
 
     if st.button("Start Training"):
         round_losses = []
-        round_accuracies = []  # Store accuracy values for each round
+        round_accuracies = []
+
         for round_num in range(1, NUM_ROUNDS + 1):
             st.write(f"## Round {round_num}")
 
             st.write("### Training Metrics for Each Client")
+            client_losses = []
+            client_accuracies = []
+
             for client in range(1, NUM_CLIENTS + 1):
-                client_loss, client_accuracy = test(net, testloader)  # Placeholder for actual client metrics
-                st.write(f"Client {client}: Loss: {client_loss}, Accuracy: {client_accuracy}")
+                train_subset, _ = random_split(trainloader.dataset, [len(trainloader.dataset) // NUM_CLIENTS] * NUM_CLIENTS)
+                trainloader_client = DataLoader(train_subset, shuffle=True, batch_size=32, collate_fn=trainloader.collate_fn)
+                train(net, trainloader_client, epochs=1)
+                client_loss, client_accuracy = test(net, testloader)
+                st.write(f"Client {client}: Loss: {client_loss:.4f}, Accuracy: {client_accuracy:.4f}")
+                client_losses.append(client_loss)
+                client_accuracies.append(client_accuracy)
+
+            avg_client_loss = sum(client_losses) / NUM_CLIENTS
+            avg_client_accuracy = sum(client_accuracies) / NUM_CLIENTS
+
+            st.write("### Average Metrics Across All Clients")
+            st.write(f"Average Loss: {avg_client_loss:.4f}, Average Accuracy: {avg_client_accuracy:.4f}")
+
+            round_losses.append(avg_client_loss)
+            round_accuracies.append(avg_client_accuracy)
 
             st.write("### Accuracy Over Rounds")
-            round_accuracies.append(client_accuracy)  # Append the accuracy for this round
-            plt.plot(range(1, round_num + 1), round_accuracies, marker='o')  # Plot accuracy over rounds
+            plt.plot(range(1, round_num + 1), round_accuracies, marker='o', label="Accuracy")
             plt.xlabel("Round")
             plt.ylabel("Accuracy")
             plt.title("Accuracy Over Rounds")
             st.pyplot()
 
             st.write("### Loss Over Rounds")
-            loss_value = random.random()  # Placeholder for loss values
-            round_losses.append(loss_value)
-            rounds = list(range(1, round_num + 1))
-            plt.plot(rounds, round_losses)
+            plt.plot(range(1, round_num + 1), round_losses, marker='o', color='red', label="Loss")
             plt.xlabel("Round")
             plt.ylabel("Loss")
             plt.title("Loss Over Rounds")
@@ -130,6 +125,141 @@ def main():
 if __name__ == "__main__":
     main()
 
+##ORIGINAL###
+
+
+# # %%writefile app.py
+
+# import streamlit as st
+# import matplotlib.pyplot as plt
+# import torch
+# from transformers import AutoTokenizer, DataCollatorWithPadding, AutoModelForSequenceClassification, AdamW
+# from datasets import load_dataset
+# from evaluate import load as load_metric
+# from torch.utils.data import DataLoader
+# import random
+
+# DEVICE = torch.device("cpu")
+# NUM_ROUNDS = 3
+
+
+
+# def load_data(dataset_name):
+#     raw_datasets = load_dataset(dataset_name)
+#     raw_datasets = raw_datasets.shuffle(seed=42)
+#     del raw_datasets["unsupervised"]
+
+#     tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+
+#     def tokenize_function(examples):
+#         return tokenizer(examples["text"], truncation=True)
+
+#     train_population = random.sample(range(len(raw_datasets["train"])), 20)
+#     test_population = random.sample(range(len(raw_datasets["test"])), 20)
+
+#     tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
+#     tokenized_datasets["train"] = tokenized_datasets["train"].select(train_population)
+#     tokenized_datasets["test"] = tokenized_datasets["test"].select(test_population)
+
+#     tokenized_datasets = tokenized_datasets.remove_columns("text")
+#     tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
+
+#     data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
+#     trainloader = DataLoader(tokenized_datasets["train"], shuffle=True, batch_size=32, collate_fn=data_collator)
+#     testloader = DataLoader(tokenized_datasets["test"], batch_size=32, collate_fn=data_collator)
+
+#     return trainloader, testloader
+
+# def train(net, trainloader, epochs):
+#     optimizer = AdamW(net.parameters(), lr=5e-5)
+#     net.train()
+#     for _ in range(epochs):
+#         for batch in trainloader:
+#             batch = {k: v.to(DEVICE) for k, v in batch.items()}
+#             outputs = net(**batch)
+#             loss = outputs.loss
+#             loss.backward()
+#             optimizer.step()
+#             optimizer.zero_grad()
+
+# def test(net, testloader):
+#     metric = load_metric("accuracy")
+#     loss = 0
+#     net.eval()
+#     for batch in testloader:
+#         batch = {k: v.to(DEVICE) for k, v in batch.items()}
+#         with torch.no_grad():
+#             outputs = net(**batch)
+#         logits = outputs.logits
+#         loss += outputs.loss.item()
+#         predictions = torch.argmax(logits, dim=-1)
+#         metric.add_batch(predictions=predictions, references=batch["labels"])
+#     loss /= len(testloader.dataset)
+#     accuracy = metric.compute()["accuracy"]
+#     return loss, accuracy
+
+
+
+
+
+# from transformers import Wav2Vec2Processor, HubertForSequenceClassification
+# import torch
+
+    
+# def main():
+#     st.write("## Federated Learning with dynamic models and datasets for mobile devices")
+#     dataset_name = st.selectbox("Dataset", ["imdb","audio_instruction_task", "amazon_polarity", "ag_news"])
+#     model_name = st.selectbox("Model", ["bert-base-uncased","facebook/hubert-base-ls960", "distilbert-base-uncased"])
+
+#     net = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2).to(DEVICE)
+#     # processor = Wav2Vec2Processor.from_pretrained(model_name)
+#     # net = HubertForSequenceClassification.from_pretrained(model_name, num_labels=2).to(DEVICE)
+
+#     # feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(model_name)
+#     # net = HubertForSequenceClassification.from_pretrained(model_name, num_labels=2).to(DEVICE)
+
+#     NUM_CLIENTS = st.slider("Number of Clients", min_value=1, max_value=10, value=2)
+#     NUM_ROUNDS = st.slider("Number of Rounds", min_value=1, max_value=10, value=3)
+
+#     trainloader, testloader = load_data(dataset_name)
+
+#     if st.button("Start Training"):
+#         round_losses = []
+#         round_accuracies = []  # Store accuracy values for each round
+#         for round_num in range(1, NUM_ROUNDS + 1):
+#             st.write(f"## Round {round_num}")
+
+#             st.write("### Training Metrics for Each Client")
+#             for client in range(1, NUM_CLIENTS + 1):
+#                 client_loss, client_accuracy = test(net, testloader)  # Placeholder for actual client metrics
+#                 st.write(f"Client {client}: Loss: {client_loss}, Accuracy: {client_accuracy}")
+
+#             st.write("### Accuracy Over Rounds")
+#             round_accuracies.append(client_accuracy)  # Append the accuracy for this round
+#             plt.plot(range(1, round_num + 1), round_accuracies, marker='o')  # Plot accuracy over rounds
+#             plt.xlabel("Round")
+#             plt.ylabel("Accuracy")
+#             plt.title("Accuracy Over Rounds")
+#             st.pyplot()
+
+#             st.write("### Loss Over Rounds")
+#             loss_value = random.random()  # Placeholder for loss values
+#             round_losses.append(loss_value)
+#             rounds = list(range(1, round_num + 1))
+#             plt.plot(rounds, round_losses)
+#             plt.xlabel("Round")
+#             plt.ylabel("Loss")
+#             plt.title("Loss Over Rounds")
+#             st.pyplot()
+
+#             st.success(f"Round {round_num} completed successfully!")
+
+#     else:
+#         st.write("Click the 'Start Training' button to start the training process.")
+
+# if __name__ == "__main__":
+#     main()
+###ORIGINAL##