Update generate

handler.py

@@ -23,6 +23,11 @@ class EndpointHandler:
         }
         self.instruction = """Extract the name of the person, the location, the hotel name and the desired date from the following hotel request"""
 
+        if torch.cuda.is_available():
+            self.device = "cuda"
+        else:
+            self.device = "cpu"
+
     def generate_prompt(
         self,
         template: str,
@@ -55,6 +60,7 @@ class EndpointHandler:
         parameters = data.pop("parameters", None)
 
         inputs = self.generate_prompt(self.template, self.instruction, inputs)
+        input_ids = input_ids.to(self.device)
         # preprocess
         self.tokenizer.pad_token_id = (
             0  # unk. we want this to be different from the eos token
@@ -65,7 +71,7 @@ class EndpointHandler:
         if parameters is not None:
             outputs = self.model.generate(input_ids, **parameters)
         else:
-            outputs = self.model.generate(input_ids)
+            outputs = self.model.generate(input_ids, max_new_tokens=64)
 
         # postprocess the prediction
         prediction = self.tokenizer.decode(outputs[0]) #, skip_special_tokens=True)