philschmid
/

gpt-j-6B-fp16-sharded

endpoints-template

Inference Endpoints

Model card Files Files and versions Community

philschmid HF staff commited on May 6, 2023

Commit

f7130c4

•

1 Parent(s): 4d500b9

Update handler.py

Files changed (1) hide show

handler.py +1 -1

handler.py CHANGED Viewed

@@ -12,7 +12,7 @@ class EndpointHandler:
         tokenizer = AutoTokenizer.from_pretrained(path)
         model = AutoModelForCausalLM.from_pretrained(path, device_map="auto",torch_dtype=dtype)
         # create inference pipeline
-        self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         inputs = data.pop("inputs", data)

         tokenizer = AutoTokenizer.from_pretrained(path)
         model = AutoModelForCausalLM.from_pretrained(path, device_map="auto",torch_dtype=dtype)
         # create inference pipeline
+        self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         inputs = data.pop("inputs", data)