Refactored handler

handler.py

@@ -1,35 +1,27 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
-from transformers import GenerationConfig
+from transformers import GenerationConfig, pipeline
 import torch
 from typing import Any, Dict
 
+dtype = torch.bfloat16 if torch.cuda.get_device_capability()[0] == 8 else torch.float16
+
 
 class EndpointHandler:
     def __init__(self, path=""):
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
-        self.model = AutoModelForCausalLM.from_pretrained(path,
-                                                          torch_dtype=torch.float16,
-                                                          device_map="auto")
+        tokenizer = AutoTokenizer.from_pretrained(path)
+        model = AutoModelForCausalLM.from_pretrained(path, device_map="auto", torch_dtype=dtype)
 
-        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
     def __call__(self, data: Dict[str, Any]) -> [str]:
-        input_text = data.pop("inputs", data)
+        inputs = data.pop("inputs", data)
         generation_config = GenerationConfig(
+            max_length=1024,
             max_new_tokens=250, do_sample=True, top_k=50,
-            eos_token_id=self.model.config.eos_token_id,
             temperature=0.8, pad_token_id=2, num_return_sequences=1,
             min_new_tokens=30, repetition_penalty=1.2,
         )
-        
-        self.model.generation_config = generation_config
-        inputs = self.tokenizer(input_text, return_tensors="pt")
-        inputs = {key: val.to(self.device) for key, val in inputs.items()}
-        outputs = self.model.generate(**inputs)
-
-        decoded_output = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-        # remove the inputs from outputs
-        decoded_output = decoded_output.replace(input_text + ' Expert: ', '')
+        output = self.pipeline(inputs, **generation_config.to_dict())
 
-        return [{'generated_text': decoded_output}]
+        return output