Update app.py

app.py

@@ -392,11 +392,11 @@
 # if __name__ == "__main__":
 #     main()
 
-
 import streamlit as st
 import matplotlib.pyplot as plt
 import torch
-from transformers import AutoTokenizer, DataCollatorWithPadding, AutoModelForSequenceClassification, AdamW
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, AdamW
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 from datasets import load_dataset, Dataset
 from evaluate import load as load_metric
 from torch.utils.data import DataLoader
@@ -413,35 +413,39 @@ import plotly.graph_objects as go
 DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 fl.common.logger.configure(identifier="myFlowerExperiment", filename="./log.txt")
 
-class CustomDataCollator(DataCollatorWithPadding):
+class CustomDataCollator:
+    def __init__(self, pad_token_id=0):
+        self.pad_token_id = pad_token_id
+
     def __call__(self, features):
-        if 'input_ids' in features[0] and isinstance(features[0]['input_ids'][0], int):
-            # Handle byte encoding case
-            max_length = max(len(f["input_ids"]) for f in features)
-            for f in features:
-                f['input_ids'] += [0] * (max_length - len(f['input_ids']))
-        return super().__call__(features)
-
-def load_data(dataset_name, train_size=20, test_size=20, num_clients=2, use_utf8=False):
+        max_length = max(len(f["input_ids"]) for f in features)
+        for f in features:
+            f['input_ids'] += [self.pad_token_id] * (max_length - len(f['input_ids']))
+        batch = {k: torch.tensor([f[k] for f in features]) for k in features[0].keys()}
+        return batch
+
+def load_data(dataset_name, train_size=20, test_size=20, num_clients=2, use_utf8=False, model_name="bert-base-uncased"):
     raw_datasets = load_dataset(dataset_name)
     raw_datasets = raw_datasets.shuffle(seed=42)
     del raw_datasets["unsupervised"]
 
-    if not use_utf8:
-        tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+    if model_name == "google/byt5-small":
+        tokenizer = T5Tokenizer.from_pretrained(model_name)
 
-        def tokenize_function(examples):
-            return tokenizer(examples["text"], truncation=True)
+        def utf8_encode_function(examples):
+            examples["input_ids"] = [tokenizer(text.encode('utf-8'), return_tensors="pt")["input_ids"].squeeze().tolist() for text in examples["text"]]
+            return examples
 
-        tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
+        tokenized_datasets = raw_datasets.map(utf8_encode_function, batched=True)
         tokenized_datasets = tokenized_datasets.remove_columns("text")
         tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
     else:
-        def utf8_encode_function(examples):
-            examples["input_ids"] = [list(text.encode('utf-8')) for text in examples["text"]]
-            return examples
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-        tokenized_datasets = raw_datasets.map(utf8_encode_function, batched=True)
+        def tokenize_function(examples):
+            return tokenizer(examples["text"], truncation=True)
+
+        tokenized_datasets = raw_datasets.map(tokenize_function, batched=True)
         tokenized_datasets = tokenized_datasets.remove_columns("text")
         tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
 
@@ -454,7 +458,7 @@ def load_data(dataset_name, train_size=20, test_size=20, num_clients=2, use_utf8
         train_datasets.append(train_dataset)
         test_datasets.append(test_dataset)
 
-    data_collator = CustomDataCollator(tokenizer=AutoTokenizer.from_pretrained("bert-base-uncased"))
+    data_collator = CustomDataCollator(pad_token_id=tokenizer.pad_token_id)
 
     return train_datasets, test_datasets, data_collator, raw_datasets
 
@@ -634,15 +638,11 @@ def read_log_file2():
 def main():
     st.write("## Federated Learning with Dynamic Models and Datasets for Mobile Devices")
     logs = read_log_file2()
-    # cleanLogs = # Define a pattern to match relevant log entries
     pattern = re.compile(r"memory|loss|accuracy|round|client", re.IGNORECASE)
     
-
-    # Filter the log data
     filtered_logs = [line for line in logs.splitlines() if pattern.search(line)]
     st.markdown(filtered_logs)
 
-    # Provide a download button for the logs
     st.download_button(
         label="Download Logs",
         data="\n".join(filtered_logs),
@@ -650,13 +650,13 @@ def main():
         mime="text/plain"
     )
     dataset_name = st.selectbox("Dataset", ["imdb", "amazon_polarity", "ag_news"])
-    model_name = st.selectbox("Model", ["bert-base-uncased", "facebook/hubert-base-ls960", "distilbert-base-uncased"])
+    model_name = st.selectbox("Model", ["bert-base-uncased", "facebook/hubert-base-ls960", "distilbert-base-uncased", "google/byt5-small"])
 
     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)
     use_utf8 = st.checkbox("Train on Byte UTF-8 Dataset", value=False)
 
-    train_datasets, test_datasets, data_collator, raw_datasets = load_data(dataset_name, num_clients=NUM_CLIENTS, use_utf8=use_utf8)
+    train_datasets, test_datasets, data_collator, raw_datasets = load_data(dataset_name, num_clients=NUM_CLIENTS, use_utf8=use_utf8, model_name=model_name)
 
     trainloaders = []
     testloaders = []
@@ -684,9 +684,6 @@ def main():
         trainloader = DataLoader(edited_train_dataset, shuffle=True, batch_size=32, collate_fn=data_collator)
         testloader = DataLoader(edited_test_dataset, batch_size=32, collate_fn=data_collator)
 
-        trainloaders.append(trainloader)
-        testloaders.append(testloader)
-
         net = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2).to(DEVICE)
         client = CustomClient(net, trainloader, testloader, client_id=i+1)
         clients.append(client)
@@ -711,18 +708,10 @@ def main():
             st.write(f"### Round {round_num + 1} ✅")
 
             logs = read_log_file2()
-            filtered_log_list = [line for line in logs.splitlines if pattern.search(line)]
+            filtered_log_list = [line for line in logs.splitlines() if pattern.search(line)]
             filtered_logs = "\n".join(filtered_log_list)
             
             st.markdown(filtered_logs)
-            # Provide a download button for the logs
-            # st.download_button(
-            #     label="Download Logs",
-            #     data=logs,
-            #     file_name="./log.txt",
-            #     mime="text/plain"
-            # )
-            # # Extract relevant data
             accuracy_pattern = re.compile(r"'accuracy': \{(\d+), ([\d.]+)\}")
             loss_pattern = re.compile(r"'loss': \{(\d+), ([\d.]+)\}")
 
@@ -733,21 +722,17 @@ def main():
             accuracies = [float(match[1]) for match in accuracy_matches]
             losses = [float(match[1]) for match in loss_matches]
 
-            # Create accuracy plot
             accuracy_fig = go.Figure()
             accuracy_fig.add_trace(go.Scatter(x=rounds, y=accuracies, mode='lines+markers', name='Accuracy'))
             accuracy_fig.update_layout(title='Accuracy over Rounds', xaxis_title='Round', yaxis_title='Accuracy')
 
-            # Create loss plot
             loss_fig = go.Figure()
             loss_fig.add_trace(go.Scatter(x=rounds, y=losses, mode='lines+markers', name='Loss'))
             loss_fig.update_layout(title='Loss over Rounds', xaxis_title='Round', yaxis_title='Loss')
 
-            # Display plots in Streamlit
             st.plotly_chart(accuracy_fig)
             st.plotly_chart(loss_fig)
 
-            # Display data table
             data = {
                 'Round': rounds,
                 'Accuracy': accuracies,
@@ -775,7 +760,6 @@ def main():
 
         st.success("Training completed successfully!")
 
-        # Display final metrics
         st.write("## Final Client Metrics")
         for client in clients:
             st.write(f"### Client {client.client_id}")
@@ -788,7 +772,6 @@ def main():
 
             st.write(" ")
 
-        # Display log.txt content
         st.write("## Training Log")
         st.write(read_log_file2())