Is there a way to set the TIKTOKEN_CACHE_DIR environment variable when using mlflow.evaluate()?

[kimminw00]

I'm currently developing in an air-gapped network environment. When using mlflow.evaluate(), I've noticed that it attempts to download the TikToken cache. To work around this, I've set the TIKTOKEN_CACHE_DIR environment variable to a local path containing the cache.

However, I've encountered an issue with the following code:

mlflow/mlflow/metrics/metric_definitions.py

Lines 70 to 75 in d88278d

	def _token_count_eval_fn(predictions, targets=None, metrics=None):
	import tiktoken
	# ref: https://github.com/openai/tiktoken/issues/75
	os.environ["TIKTOKEN_CACHE_DIR"] = ""
	encoding = tiktoken.get_encoding("cl100k_base")

The line os.environ["TIKTOKEN_CACHE_DIR"] = "" resets the TIKTOKEN_CACHE_DIR environment variable to an empty string, which causes the cache to be downloaded again instead of using the local cache.

Is there a way to override this line, so that the TIKTOKEN_CACHE_DIR environment variable is respected and the local cache is used instead of downloading it again?

I'd appreciate any help or guidance on how to resolve this issue. Thank you!

[omorlui]

We need that too!

[B-Step62]

Thanks for the suggestion! Filed a PR to allow overriding cache dir and only set empty string when it is not set by users.
#16303