update

utils/model.py

@@ -23,16 +23,17 @@ class Model(torch.nn.Module):
         
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
         self.name = model_name
+
         logging.info(f'start loading model {self.name}')
         self.pipeline = transformers.pipeline(
-            "summarization",
+            "summarization" if model_name=="google-t5/t5-large" else "text-generation",
             model=model_name,
             tokenizer=self.tokenizer,
             torch_dtype=torch.bfloat16,
             device_map="auto",
         )
-
         logging.info(f'Loaded model {self.name}')
+
         self.update()
 
     @classmethod
@@ -58,6 +59,7 @@ class Model(torch.nn.Module):
                 num_return_sequences=1,
                 eos_token_id=self.tokenizer.eos_token_id,
             )
+            return sequences[-1]['summary_text']
         else:
             sequences = self.pipeline(
                 content,
@@ -68,5 +70,4 @@ class Model(torch.nn.Module):
                 eos_token_id=self.tokenizer.eos_token_id,
                 return_full_text=False
             )
-        
-        return sequences[-1]['summary_text']
+            return sequences[-1]['generated_text']