Update app.py

app.py

@@ -391,11 +391,12 @@
 
 # if __name__ == "__main__":
 #     main()
+
 import streamlit as st
 import matplotlib.pyplot as plt
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, AdamW
-from transformers import T5Tokenizer, T5ForConditionalGeneration
+from transformers import ByT5Tokenizer, ByT5ForConditionalGeneration
 from datasets import load_dataset, Dataset
 from evaluate import load as load_metric
 from torch.utils.data import DataLoader
@@ -429,11 +430,11 @@ def load_data(dataset_name, train_size=20, test_size=20, num_clients=2, use_utf8
     del raw_datasets["unsupervised"]
 
     if model_name == "google/byt5-small":
-        tokenizer = T5Tokenizer.from_pretrained(model_name)
+        tokenizer = ByT5Tokenizer.from_pretrained(model_name)
 
         def utf8_encode_function(examples):
-            encoded_texts = [text.encode('utf-8') for text in examples["text"]]
-            examples["input_ids"] = [tokenizer(list(encoded_text), return_tensors="pt", padding='max_length', truncation=True, max_length=512)["input_ids"].squeeze().tolist() for encoded_text in encoded_texts]
+            encoded_texts = [list(text.encode('utf-8')) for text in examples["text"]]
+            examples["input_ids"] = [tokenizer(text, return_tensors="pt", padding='max_length', truncation=True, max_length=512)["input_ids"].squeeze().tolist() for text in encoded_texts]
             return examples
 
         tokenized_datasets = raw_datasets.map(utf8_encode_function, batched=True)
@@ -683,7 +684,11 @@ 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)
 
-        net = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2).to(DEVICE)
+        if model_name == "google/byt5-small":
+            net = ByT5ForConditionalGeneration.from_pretrained(model_name).to(DEVICE)
+        else:
+            net = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2).to(DEVICE)
+
         client = CustomClient(net, trainloader, testloader, client_id=i+1)
         clients.append(client)
 
@@ -786,3 +791,4 @@ def main():
 if __name__ == "__main__":
     main()
 
+