Update app.py

app.py

@@ -391,12 +391,11 @@
 
 # if __name__ == "__main__":
 #     main()
-
 import streamlit as st
 import matplotlib.pyplot as plt
 import torch
 from transformers import AutoTokenizer, AutoModelForSequenceClassification, AdamW
-from transformers import ByT5Tokenizer, ByT5ForConditionalGeneration
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 from datasets import load_dataset, Dataset
 from evaluate import load as load_metric
 from torch.utils.data import DataLoader
@@ -430,7 +429,7 @@ 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 = ByT5Tokenizer.from_pretrained(model_name)
+        tokenizer = T5Tokenizer.from_pretrained(model_name)
 
         def utf8_encode_function(examples):
             encoded_texts = [list(text.encode('utf-8')) for text in examples["text"]]
@@ -685,7 +684,7 @@ def main():
         testloader = DataLoader(edited_test_dataset, batch_size=32, collate_fn=data_collator)
 
         if model_name == "google/byt5-small":
-            net = ByT5ForConditionalGeneration.from_pretrained(model_name).to(DEVICE)
+            net = T5ForConditionalGeneration.from_pretrained(model_name).to(DEVICE)
         else:
             net = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2).to(DEVICE)
 
@@ -790,5 +789,3 @@ def main():
 
 if __name__ == "__main__":
     main()
-
-