IVYXSTUDIO · IVY NODE

Tokenizer Model Loader

ivy.node.tokenizer-model-loader · v0.0.0

ivyx✓

Loads a pre-trained tokenizer and model from Hugging Face for text embedding generation. This node downloads and initializes both the tokenizer and model required for text vectorization. The outputs are used with text-vectorizer node to create embedding instances. Requires a Hugging Face access token for authentication (especially for private models). This is typically the first step in ML pipelines that require text embeddings.

#Tokenizer#Model#Loader

Inputs

Field	Type	Description
hugging_face_tokenrequired	string	This token is Hugging Face access token
model_namerequired	string	The name of the model to load
features	object	Features
labels	object	Labels

Outputs

Field	Type	Description
tokenizerrequired	object	Custom type: PreTrainedTokenizer
modelrequired	object	Custom type: PreTrainedModel
features	object	Features
labels	object	Labels

Source

python

# Input preparation
inp = __ivy_ctx__["nodes"][__ivy_node_id__]["input"]
hugging_face_token = inp["hugging_face_token"]
model_name = inp["model_name"]

# Compute
from transformers import AutoTokenizer, AutoModel, PreTrainedTokenizer, PreTrainedModel
import os

class TokenizerModelLoader:
    """
    A class to load a tokenizer and model for generating embeddings.
    """

    def __init__(self, model_name: str):
        """
        Initialize the TokenizerModelLoader with a model name.

        :param model_name: The name of the model to load.
        """
        self.model_name = model_name
        self.tokenizer = None
        self.model = None

    def load_tokenizer(self) -> PreTrainedTokenizer:
        """
        Load the tokenizer for the specified model.

        :return: The tokenizer.
        :raises: RuntimeError if the tokenizer fails to load.
        """
        try:
            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
            print(f"Tokenizer loaded successfully for model: {self.model_name}")
            return self.tokenizer
        except Exception as e:
            print(f"Failed to load tokenizer for model {self.model_name}: {e}")
            raise RuntimeError(f"Tokenizer loading failed: {e}")

    def load_model(self) -> PreTrainedModel:
        """
        Load the model for the specified model name.

        :return: The model.
        :raises: RuntimeError if the model fails to load.
        """
        try:
            self.model = AutoModel.from_pretrained(self.model_name)
            print(f"Model loaded successfully for model: {self.model_name}")
            return self.model
        except Exception as e:
            print(f"Failed to load model {self.model_name}: {e}")
            raise RuntimeError(f"Model loading failed: {e}")

os.environ["HF_TOKEN"] = hugging_face_token

tokenizer_model_loader = TokenizerModelLoader(model_name)

tokenizer = tokenizer_model_loader.load_tokenizer()

model = tokenizer_model_loader.load_model()

# Output collection (runner reads __ivy_ctx__)
out = __ivy_ctx__["nodes"][__ivy_node_id__]["output"]
out["tokenizer"] = tokenizer
out["model"] = model

Tests

Requires: python:3.11

basic-model-loading
Load tokenizer and model with valid inputs (should succeed).
model-loadingtokenizerbasic
error-invalid-model
Invalid model name should handle error gracefully.
error-handlinginvalid-input