adding temperature hyper params

handler.py

@@ -37,23 +37,29 @@ class EndpointHandler():
           
 
     def _secure_inputs(self, data: Dict[str, Any]):
-        inputs = data.get("inputs", None)
-        if inputs is None:
-            return [{"error": "inputs should be shaped like {'inputs': <string or List of strings (abstracts)>}"}], False
+        if not isinstance(data, dict):
+          return [{"error": "inputs should be shaped like {'temperature': float, 'inputs': <string or List of strings (abstracts)>}"}], False
+
+        if not 'inputs' in data:
+          return [{"error": "inputs should be shaped like {'temperature': float, 'inputs': <string or List of strings (abstracts)>}"}], False
+
+        temperature = data.get("temperature", 0.01)
+        inputs = data["inputs"]
 
         if isinstance(inputs, str):
             inputs = [inputs]
 
-        return inputs, True
+        return inputs, temperature, True
 
+    
     def _format_inputs(self, inputs: list[str]):
         prompts = [self.summary_prompt.format(abstract, "") for abstract in inputs]
         prompts_lengths = [len(prompt) for prompt in prompts]
         return prompts, prompts_lengths
 
-    def _generate_outputs(self, inputs):
+    def _generate_outputs(self, inputs, temperature):
         tokenized = self.tokenizer(inputs, return_tensors="pt", padding=True).to("cuda")
-        outputs = self.model.generate(**tokenized, max_new_tokens=500, use_cache=True)
+        outputs = self.model.generate(**tokenized, temperature=temperature, max_new_tokens=500, use_cache=True)
         decoded = self.tokenizer.batch_decode(outputs, skip_special_tokens=True) 
         return decoded
 
@@ -67,13 +73,13 @@ class EndpointHandler():
         
         
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
-        inputs, is_secure = self._secure_inputs(data)
+        inputs, temperature, is_secure = self._secure_inputs(data)
         
         if not is_secure:
             return inputs
         
         inputs, inputs_length = self._format_inputs(inputs)
-        outputs = self._generate_outputs(inputs)
+        outputs = self._generate_outputs(inputs, temperature)
         outputs = self._format_outputs(outputs, inputs_length)
 
         outputs = [{"summary": output_} for output_ in outputs]