switch to AutoModelForCausalLM

handler.py

@@ -1,21 +1,21 @@
 import torch
 from typing import Dict, List, Any
-from transformers import AutoModel, AutoTokenizer, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 # check for GPU
 device = 0 if torch.cuda.is_available() else -1
 
+
 class EndpointHandler:
     def __init__(self, path=""):
         # load the model
         tokenizer = AutoTokenizer.from_pretrained(path)
-        model = AutoModel.from_pretrained(path, low_cpu_mem_usage=True)
+        # 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)
         # create inference pipeline
-        # Do I have to check device?
-        self.pipeline = pipeline(
-            "text-generation", model=model, tokenizer=tokenizer, device=device)
+        self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
 
-    # (Might have to adjust typing)
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", None)

test.py

@@ -1,13 +0,0 @@
-from handler import EndpointHandler
-
-# init handler
-my_handler = EndpointHandler(path=".")
-
-# prepare sample payload
-payload = {"inputs": "I am Bob and I want to "}
-
-# test the handler
-pred=my_handler(payload)
-
-# show results
-print("pred", pred)