models_2.py

import tensorflow as tf
from config_utils import tf_gpu
tf_gpu()
tf.experimental.numpy.experimental_enable_numpy_behavior(prefer_float32=True)

from huggingface_hub import from_pretrained_keras
from datasets import load_dataset

data = load_dataset("jkwiatkowski/raven", split="val")
model = from_pretrained_keras("jkwiatkowski/raven")

def convert(data):
    return {
        'inputs': tf.cast(data['inputs'], dtype="uint8"),
        'index': tf.cast(data['index'], dtype="uint8"),
        'target': tf.cast(data['target'], dtype="int8"),
    }

print("xD")