Merge pull request #13 from Shivanandroy/byt5

Byt5 support, transformers upgrade, dropping onnx support

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

setup.py

@@ -8,7 +8,7 @@
 
 setuptools.setup(
     name="simplet5",
-    version="0.1.2",
+    version="0.1.3",
     license="apache-2.0",
     author="Shivanand Roy",
     author_email="[email protected]",
@@ -41,9 +41,10 @@
     install_requires=[
         "sentencepiece",
         "torch>=1.7.0,!=1.8.0",  # excludes torch v1.8.0
-        "transformers==4.6.1",
-        "pytorch-lightning==1.3.3",
-        "fastt5==0.0.6",
+        "transformers==4.10.0",
+        "pytorch-lightning==1.4.5",
+        "tqdm"
+        # "fastt5==0.0.7",
     ],
     classifiers=[
         "Intended Audience :: Developers",

simplet5/simplet5.py

@@ -6,12 +6,14 @@
     AdamW,
     T5ForConditionalGeneration,
     MT5ForConditionalGeneration,
+    ByT5Tokenizer,
     PreTrainedTokenizer,
     T5TokenizerFast as T5Tokenizer,
     MT5TokenizerFast as MT5Tokenizer,
 )
 from transformers import AutoTokenizer
-from fastT5 import export_and_get_onnx_model
+
+# from fastT5 import export_and_get_onnx_model
 from torch.utils.data import Dataset, DataLoader
 from transformers import AutoModelWithLMHead, AutoTokenizer
 import pytorch_lightning as pl
@@ -246,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)
 
@@ -282,11 +284,11 @@ def from_pretrained(self, model_type="t5", model_name="t5-base") -> None:
             self.model = MT5ForConditionalGeneration.from_pretrained(
                 f"{model_name}", return_dict=True
             )
-        # elif model_type == "byt5":
-        #     self.tokenizer = ByT5Tokenizer.from_pretrained(f"{model_name}")
-        #     self.model = T5ForConditionalGeneration.from_pretrained(
-        #         f"{model_name}", return_dict=True
-        #     )
+        elif model_type == "byt5":
+            self.tokenizer = ByT5Tokenizer.from_pretrained(f"{model_name}")
+            self.model = T5ForConditionalGeneration.from_pretrained(
+                f"{model_name}", return_dict=True
+            )
 
     def train(
         self,
@@ -385,9 +387,9 @@ def load_model(
         elif model_type == "mt5":
             self.model = MT5ForConditionalGeneration.from_pretrained(f"{model_dir}")
             self.tokenizer = MT5Tokenizer.from_pretrained(f"{model_dir}")
-        # elif model_type == "byt5":
-        #     self.model = T5ForConditionalGeneration.from_pretrained(f"{model_dir}")
-        #     self.tokenizer = ByT5Tokenizer.from_pretrained(f"{model_dir}")
+        elif model_type == "byt5":
+            self.model = T5ForConditionalGeneration.from_pretrained(f"{model_dir}")
+            self.tokenizer = ByT5Tokenizer.from_pretrained(f"{model_dir}")
 
         if use_gpu:
             if torch.cuda.is_available():
@@ -459,18 +461,18 @@ def predict(
         ]
         return preds
 
-    def convert_and_load_onnx_model(self, model_dir: str):
-        """ returns ONNX model """
-        self.onnx_model = export_and_get_onnx_model(model_dir)
-        self.onnx_tokenizer = AutoTokenizer.from_pretrained(model_dir)
-
-    def onnx_predict(self, source_text: str):
-        """ generates prediction from ONNX model """
-        token = self.onnx_tokenizer(source_text, return_tensors="pt")
-        tokens = self.onnx_model.generate(
-            input_ids=token["input_ids"],
-            attention_mask=token["attention_mask"],
-            num_beams=2,
-        )
-        output = self.onnx_tokenizer.decode(tokens.squeeze(), skip_special_tokens=True)
-        return output
+    # def convert_and_load_onnx_model(self, model_dir: str):
+    #     """ returns ONNX model """
+    #     self.onnx_model = export_and_get_onnx_model(model_dir)
+    #     self.onnx_tokenizer = AutoTokenizer.from_pretrained(model_dir)
+
+    # def onnx_predict(self, source_text: str):
+    #     """ generates prediction from ONNX model """
+    #     token = self.onnx_tokenizer(source_text, return_tensors="pt")
+    #     tokens = self.onnx_model.generate(
+    #         input_ids=token["input_ids"],
+    #         attention_mask=token["attention_mask"],
+    #         num_beams=2,
+    #     )
+    #     output = self.onnx_tokenizer.decode(tokens.squeeze(), skip_special_tokens=True)
+    #     return output