Fine tunning for zero shot on tensorflow

Has anyone seen a code snippet for fine tuning a model on NLI using the transformer library on a tensorflow backbone?

Specifically:

I have a data set with specialized vocabulary, for which I want to apply zero shot classification. I first fine tuned on a maskedLM task but now the model has lost the NLI heads. From my understanding fine tuning my new model on an NLI task next should achieve the desired result.

I’ve seen a PyTorch implementation for NLI tunning, but I’m not as familiar with PyTorch.

I’m trying to fine tune my model on the multi_nli dataset.

Any help is appreciated, including alternative approaches.