Upload T5.py

T5.py

@@ -0,0 +1,119 @@
+
+from transformers import T5TokenizerFast, T5ForConditionalGeneration
+from torch import Tensor
+from torch.nn import Module
+from typing import List, Optional, Tuple
+import torch, os
+import torch.nn.functional as F
+
+
+class T5(Module):
+    '''
+    T5 model from: https://huggingface.co/docs/transformers/model_doc/t5
+    '''
+
+    def __init__(self, 
+                 variant:str="t5-small",
+                 max_source_length:int=256, 
+                 max_target_length:int=128,
+                 optimizer_config:dict={},
+                ):
+
+        # Assertions
+        assert variant in ["t5-small", "t5-base", "t5-large"]
+
+        super().__init__()
+
+        self.variant            = variant
+        self.max_source_length  = max_source_length
+        self.max_target_length  = max_target_length
+
+        # Tokenizer & model
+        self.tokenizer          = T5TokenizerFast.from_pretrained(self.variant,
+                                                                  model_max_length=self.max_source_length)
+        self.model              = T5ForConditionalGeneration.from_pretrained(self.variant)
+        
+        # Optimizer
+        self.optimizer          = torch.optim.AdamW(self.parameters(), **optimizer_config)
+        
+        # Scheduler
+        self.scheduler          = None
+
+
+    def tokenize(self, input:List[str]):
+
+        out = self.tokenizer(input, max_length=self.max_source_length,
+                             truncation=True, padding=True, 
+                             return_tensors="pt")
+
+        return out.input_ids.cuda(), out.attention_mask.cuda()
+
+
+    def forward(self, input:List[str], label:Optional[List[str]]=None) -> Tuple[Tensor, Optional[Tensor]]:
+
+        '''
+        Will receive input and target string and produce the final output as tensor (not decoded)
+        when target is not None, it will give the loss functions with the output as tuple
+        '''
+        
+        input_ids, input_masks = self.tokenize(input)
+        
+        if label is not None:
+            label_ids, label_masks = self.tokenize(label)    
+            output = self.model(input_ids=input_ids, labels=label_ids)
+            return output.logits, output.loss
+            
+        return self.model.generate(input_ids=input_ids,
+                                   max_new_tokens=self.max_target_length), None
+
+
+
+    def predict(self, input:List[str]) -> List[str]:
+        
+        '''
+        Will generate the target output as string
+        '''
+        
+        logits, loss = self.forward(input=input)
+        return self.tokenizer.batch_decode(logits, skip_special_tokens=True)
+        
+
+
+if __name__ == '__main__':
+
+    '''
+    Implement a tester class similar to T5-old.py to test if it works
+    '''
+    
+    model = T5('t5-small')
+    model.to('cuda')
+
+    #inputs = [
+        #"translate English to German: Thank you so much, Chris.",
+        #"translate English to German: I have been blown away by this conference, and I want to thank all of you for the many nice comments about what I had to say the other night.",
+        #"translate German to English: Es ist mir wirklich eine Ehre, zweimal auf dieser Bühne stehen zu dürfen. Tausend Dank dafür.",
+    #]
+
+    #targets = [
+        #"Vielen Dank, Chris.",
+        #"Ich bin wirklich begeistert von dieser Konferenz, und ich danke Ihnen allen für die vielen netten Kommentare zu meiner Rede vorgestern Abend.",
+        #"And it's truly a great honor to have the opportunity to come to this stage twice; I'm extremely grateful.",
+    #]
+
+    inputs = ["Good Morning, How are you?"]
+    targets = ["Buongiorno, come stai?"]
+
+    logits, loss = model.forward(inputs, targets)
+    print('Model forward')
+    print('logits: ', logits)
+    print('loss: ', loss)
+    
+    outputs = model.predict(inputs)
+    
+    #print('OUTPUT')
+    #print(outputs)
+    for (inp, out), tar in zip(zip(inputs, outputs), targets):
+        print(f"Input: \n{inp}\n\nOutput: \n{out}\n\nTarget: \n{tar}\n\n")
+
+
+