Update app.py

app.py

@@ -391,7 +391,6 @@
 
 # if __name__ == "__main__":
 #     main()
-
 import streamlit as st
 import matplotlib.pyplot as plt
 import torch
@@ -412,6 +411,15 @@ 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):
+    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):
     raw_datasets = load_dataset(dataset_name)
     raw_datasets = raw_datasets.shuffle(seed=42)
@@ -444,7 +452,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 = DataCollatorWithPadding(tokenizer=AutoTokenizer.from_pretrained("bert-base-uncased"))
+    data_collator = CustomDataCollator(tokenizer=AutoTokenizer.from_pretrained("bert-base-uncased"))
 
     return train_datasets, test_datasets, data_collator, raw_datasets
 
@@ -701,7 +709,7 @@ 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)