Update handler.py

handler.py

@@ -36,28 +36,22 @@ class EndpointHandler:
     def __call__(self, data: Any):
         inputs = data.pop("inputs", data)
         parameters = data.pop("parameters", None)
-        if parameters.get("no_generation", False):
-            input_tokens = self.tokenizer.batch_encode_plus(
+        input_tokens = self.tokenizer.batch_encode_plus(
                 [inputs], return_tensors="pt", padding=False
             )
-            for t in input_tokens:
-                if torch.is_tensor(input_tokens[t]):
-                    input_tokens[t] = input_tokens[t].to(torch.cuda.current_device())
-
-            logits = self.model(
-                input_ids=input_tokens["input_ids"],
-                attention_mask=input_tokens["attention_mask"],
-            )[0]
-            log_likelihood = self.compute_log_likelihood(
-                logits, input_tokens["input_ids"]
-            )
-            return (logits, log_likelihood)
-        if parameters is not None:
-            prediction = self.pipeline(inputs, **parameters)
-        else:
-            prediction = self.pipeline(inputs)
-        return prediction
+        for t in input_tokens:
+            if torch.is_tensor(input_tokens[t]):
+                input_tokens[t] = input_tokens[t].to(torch.cuda.current_device())
 
+        logits = self.model(
+            input_ids=input_tokens["input_ids"],
+            attention_mask=input_tokens["attention_mask"],
+        )[0]
+        log_likelihood = self.compute_log_likelihood(
+            logits, input_tokens["input_ids"]
+        )
+        return (logits, log_likelihood)
+        
 
 # if __name__ == "__main__":
 #     model = EndpointHandler("TinyLlama/TinyLlama-1.1B-Chat-v1.0")