Update app.py

app.py

@@ -1,13 +1,13 @@
 # %%writefile app.py
-# %%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 datasets import load_dataset, Dataset
 from evaluate import load as load_metric
 from torch.utils.data import DataLoader
+import pandas as pd
 import random
 import warnings
 from collections import OrderedDict
@@ -40,10 +40,7 @@ def load_data(dataset_name, train_size=20, test_size=20, num_clients=2):
 
     data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
 
-    trainloaders = [DataLoader(ds, shuffle=True, batch_size=32, collate_fn=data_collator) for ds in train_datasets]
-    testloaders = [DataLoader(ds, batch_size=32, collate_fn=data_collator) for ds in test_datasets]
-
-    return trainloaders, testloaders
+    return train_datasets, test_datasets, data_collator
 
 def train(net, trainloader, epochs):
     optimizer = AdamW(net.parameters(), lr=5e-5)
@@ -107,7 +104,30 @@ def main():
     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)
 
-    trainloaders, testloaders = load_data(dataset_name, num_clients=NUM_CLIENTS)
+    train_datasets, test_datasets, data_collator = load_data(dataset_name, num_clients=NUM_CLIENTS)
+
+    trainloaders = []
+    testloaders = []
+
+    for i in range(NUM_CLIENTS):
+        st.write(f"### Client {i+1} Datasets")
+
+        train_df = pd.DataFrame(train_datasets[i])
+        test_df = pd.DataFrame(test_datasets[i])
+
+        st.write("#### Train Dataset")
+        edited_train_df = st.experimental_data_editor(train_df, key=f"train_{i}")
+        st.write("#### Test Dataset")
+        edited_test_df = st.experimental_data_editor(test_df, key=f"test_{i}")
+
+        edited_train_dataset = Dataset.from_pandas(edited_train_df)
+        edited_test_dataset = Dataset.from_pandas(edited_test_df)
+
+        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)
 
     if st.button("Start Training"):
         round_losses = []
@@ -148,6 +168,7 @@ if __name__ == "__main__":
     main()
 
 
+
 ##ORIGINAL###