import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_text

preprocessor = hub.load("https://tfhub.dev/jeongukjae/distilkobert_cased_preprocess/1")
tokenize_with_offsets = hub.KerasLayer(preprocessor.tokenize_with_offsets)

@tf.function(input_signature=[tf.TensorSpec([None], dtype=tf.string)])
def _tokenize(x):
    tokens, starts, ends = tokenize_with_offsets(x)
    x = tf.map_fn(
        lambda x: tf.repeat(x[0], tf.size(x[1])),
        [x, tokens],
        fn_output_signature=tf.RaggedTensorSpec([None], dtype=tf.string),
    )  # 이 부분은 더 편하게 바꿀 방법이 있을 것 같다.

    zeros = tf.cast(tf.repeat(0, tf.size(starts)), starts.dtype)
    starts = tf.strings.length(x.with_flat_values(tf.strings.substr(x.values, zeros, starts.values)), unit='UTF8_CHAR')
    ends = tf.strings.length(x.with_flat_values(tf.strings.substr(x.values, zeros, ends.values)), unit='UTF8_CHAR')
    return tokens, starts, ends