Update handler.py

handler.py

@@ -8,14 +8,11 @@ DEFAULT_MAX_NEW_TOKENS = 10
 
 class EndpointHandler():
     def __init__(self, path: str = ""):
-        assert torch.cuda.device_count() >= 4, f"Only found access to {torch.cuda.device_count()} GPUs"
 
         self.tokenizer = AutoTokenizer.from_pretrained(path)
         self.model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.float16)
         self.model = self.model.to('cuda:0')
 
-        self.model.parallelize()
-
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         """
         Args:
@@ -27,14 +24,14 @@ class EndpointHandler():
         
         prompts = [f"<human>: {prompt}\n<bot>:" for prompt in data["inputs"]]
 
-        print("prompts")
-        raise ValueError(inputs)
-        
+        self.tokenizer.pad_token = self.tokenizer.eos_token
         inputs = self.tokenizer(prompts, padding=True, return_tensors='pt').to(self.model.device)
         input_length = inputs.input_ids.shape[1]
+
         outputs = self.model.generate(
-            **inputs, max_new_tokens=512, do_sample=True, temperature=0.7, top_p=0.7, top_k=50
+            **inputs, **data["parameters"]
          )
-        output_strs = self.tokenizer.batch_decode(outputs, skip_special_tokens=True)
         
-        return {"generated_text": output_strs}
+        output_strs = self.tokenizer.batch_decode(outputs[:, input_length:], skip_special_tokens=True)
+        
+        return [{"generated_text": output_strs}]