updated readme.md

Author: Shivanand Roy

README.md

@@ -1,18 +1,17 @@
 <img align="center" src="data/st5.png" alt="simpleT5">
 
 <p align="center">
-<b>Quickly train T5 models in just 3 lines of code with ONNX inference
+<b>Quickly train T5/mT5/byT5 models in just 3 lines of code
 </b>
 </p>
-
 <p align="center">
 <a href="https://badge.fury.io/py/simplet5"><img src="https://badge.fury.io/py/simplet5.svg" alt="PyPI version" height="18"></a>
 
 <a href="https://badge.fury.io/py/simplet5">
         <img alt="Stars" src="https://img.shields.io/github/stars/Shivanandroy/simpleT5?color=blue">
     </a>
 <a href="https://pepy.tech/project/simplet5">
-        <img alt="Stats" src="https://static.pepy.tech/personalized-badge/simplet5?period=month&units=international_system&left_color=black&right_color=orange&left_text=downloads/month">
+        <img alt="Stats" src="https://static.pepy.tech/personalized-badge/simplet5?period=total&units=international_system&left_color=black&right_color=brightgreen&left_text=Downloads">
     </a>
 <a href="https://opensource.org/licenses/MIT">
         <img alt="License" src="https://img.shields.io/badge/License-MIT-yellow.svg">
@@ -41,7 +40,7 @@ from simplet5 import SimpleT5
 # instantiate
 model = SimpleT5()
 
-# load (supports t5, mt5 models)
+# load (supports t5, mt5, byT5 models)
 model.from_pretrained("t5","t5-base")
 
 # train
@@ -63,10 +62,6 @@ model.load_model("t5","path/to/trained/model/directory", use_gpu=False)
 # predict
 model.predict("input text for prediction")
 
-# need faster inference on CPU, get ONNX support
-model.convert_and_load_onnx_model("path/to/T5 model/directory")
-model.onnx_predict("input text for prediction")
-
 ```
 ## Articles
 - [Geek Culture: simpleT5 — Train T5 Models in Just 3 Lines of Code](https://medium.com/geekculture/simplet5-train-t5-models-in-just-3-lines-of-code-by-shivanand-roy-2021-354df5ae46ba)

requirements.txt

@@ -1,6 +1,5 @@
 pandas
 sentencepiece
 torch>=1.7.0,!=1.8.0
-transformers==4.6.1
-pytorch-lightning==1.3.3
-fastt5==0.0.6
+transformers==4.10.0
+pytorch-lightning==1.4.5

simplet5/simplet5.py

@@ -248,7 +248,7 @@ def training_epoch_end(self, training_step_outputs):
             torch.mean(torch.stack([x["loss"] for x in training_step_outputs])).item(),
             4,
         )
-        path = f"{self.outputdir}/SimpleT5-epoch-{self.current_epoch}-train-loss-{str(avg_traning_loss)}"
+        path = f"{self.outputdir}/simplet5-epoch-{self.current_epoch}-train-loss-{str(avg_traning_loss)}"
         self.tokenizer.save_pretrained(path)
         self.model.save_pretrained(path)