Update t5_model.py

t5_model.py

@@ -1,3 +1,5 @@
+import pytorch_lightning as pl
+from transformers import AutoModelForSeq2SeqLM
 
 class T5(pl.LightningModule):
     def __init__(self, lr=5e-5, num_train_epochs=15, warmup_steps=1000):
@@ -85,6 +87,9 @@ class T5(pl.LightningModule):
         # Save the model
         self.model.push_to_hub(model_name, organization)
 
+    def from_pretrained(self, model_path):
+        AutoModelForSeq2SeqLM.from_pretrained(model_path)
+
         
     def train_dataloader(self):
         return train_dataloader