Switch to AutoModelForSeq2SeqLM

handler.py

@@ -1,6 +1,7 @@
 import torch
 from typing import Dict, List, Any
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+# from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
 # check for GPU
 device = 0 if torch.cuda.is_available() else -1
@@ -11,8 +12,8 @@ class EndpointHandler:
         # load the model
         tokenizer = AutoTokenizer.from_pretrained(path)
         # model = AutoModel.from_pretrained(path, low_cpu_mem_usage=True)
-        model = AutoModelForCausalLM.from_pretrained(path, low_cpu_mem_usage=True)
-        # model = AutoModelForSeq2SeqLM.from_pretrained(path, low_cpu_mem_usage=True)
+        # model = AutoModelForCausalLM.from_pretrained(path, low_cpu_mem_usage=True)
+        model = AutoModelForSeq2SeqLM.from_pretrained(path, low_cpu_mem_usage=True)
         # create inference pipeline
         self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)