JIT compilation failed

[yangjen]

I was trying to run the below sample code (python API) in Jupyter notebook but encounter this error

Sample Code:

import bleurt
from bleurt import score

checkpoint = "bleurt/test_checkpoint"
references = ["This is a test."]
candidates = ["This is the test."]

scorer = score.BleurtScorer(checkpoint)
scores = scorer.score(references=references, candidates=candidates)
assert isinstance(scores, list) and len(scores) == 1
print(scores)

Error:

2024-03-30 23:05:27.516741: W tensorflow/core/framework/op_kernel.cc:1827] UNKNOWN: JIT compilation failed.
2024-03-30 23:05:27.516767: W tensorflow/core/framework/local_rendezvous.cc:404] Local rendezvous is aborting with status: UNKNOWN: JIT compilation failed.
	 [[{{node bert/embeddings/LayerNorm/batchnorm/Rsqrt}}]]

JIT compilation failed. 
[[{{node bert/embeddings/LayerNorm/batchnorm/Rsqrt}}]] [Op:__inference_pruned_2804]. 

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
File ~/anaconda3/envs/nlu/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/function_type_utils.py:442, in bind_function_inputs(args, kwargs, function_type, default_values)
    441 try:
--> 442   bound_arguments = function_type.bind_with_defaults(
    443       args, sanitized_kwargs, default_values
    444   )
    445 except Exception as e:

File ~/anaconda3/envs/nlu/lib/python3.9/site-packages/tensorflow/core/function/polymorphism/function_type.py:264, in FunctionType.bind_with_defaults(self, args, kwargs, default_values)
    263 """Returns BoundArguments with default values filled in."""
--> 264 bound_arguments = self.bind(*args, **kwargs)
    265 bound_arguments.apply_defaults()

File ~/anaconda3/envs/nlu/lib/python3.9/inspect.py:3045, in Signature.bind(self, *args, **kwargs)
   3041 """Get a BoundArguments object, that maps the passed `args`
   3042 and `kwargs` to the function's signature.  Raises `TypeError`
   3043 if the passed arguments can not be bound.
   3044 """
-> 3045 return self._bind(args, kwargs)

File ~/anaconda3/envs/nlu/lib/python3.9/inspect.py:3034, in Signature._bind(self, args, kwargs, partial)
   3033     else:
-> 3034         raise TypeError(
   3035             'got an unexpected keyword argument {arg!r}'.format(
   3036                 arg=next(iter(kwargs))))
   3038 return self._bound_arguments_cls(self, arguments)

TypeError: got an unexpected keyword argument 'input_ids'

The above exception was the direct cause of the following exception:

TypeError                                 Traceback (most recent call last)
File ~/anaconda3/envs/nlu/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/concrete_function.py:1179, in ConcreteFunction._call_impl(self, args, kwargs)
   1178 try:
-> 1179   return self._call_with_structured_signature(args, kwargs)
   1180 except TypeError as structured_err:

File ~/anaconda3/envs/nlu/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/concrete_function.py:1259, in ConcreteFunction._call_with_structured_signature(self, args, kwargs)
   1245 """Executes the wrapped function with the structured signature.
   1246 
   1247 Args:
   (...)
   1256     of this `ConcreteFunction`.
   1257 """
   1258 bound_args = (
-> 1259     function_type_utils.canonicalize_function_inputs(
   1260         args, kwargs, self.function_type)
   1261 )
   1262 filtered_flat_args = self.function_type.unpack_inputs(bound_args)

File ~/anaconda3/envs/nlu/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/function_type_utils.py:422, in canonicalize_function_inputs(args, kwargs, function_type, default_values, is_pure)
    421   args, kwargs = _convert_variables_to_tensors(args, kwargs)
--> 422 bound_arguments = bind_function_inputs(
    423     args, kwargs, function_type, default_values
    424 )
    425 return bound_arguments

File ~/anaconda3/envs/nlu/lib/python3.9/site-packages/tensorflow/python/eager/polymorphic_function/function_type_utils.py:446, in bind_function_inputs(args, kwargs, function_type, default_values)
    445 except Exception as e:
--> 446   raise TypeError(
    447       f"Binding inputs to tf.function failed due to `{e}`. "
    448       f"Received args: {args} and kwargs: {sanitized_kwargs} for signature:"
    449       f" {function_type}."
    450   ) from e
    451 return bound_arguments

TypeError: Binding inputs to tf.function failed due to `got an unexpected keyword argument 'input_ids'`. Received args: () and kwargs: {'input_ids': <tf.Tensor: shape=(1, 512), dtype=int64, numpy=

.......

nknownError: Graph execution error:
Detected at node bert/embeddings/LayerNorm/batchnorm/Rsqrt defined at (most recent call last)

I followed the readme instruction to install and it worked on command-line. Any idea on this error?