import numpy as np
from ai_edge_litert.interpreter import Interpreter
from transformers import AutoTokenizer

# Load model
interpreter = Interpreter(model_path="google_mobilebert-uncased.tflite")
interpreter.allocate_tensors()

tokenizer = AutoTokenizer.from_pretrained("google/mobilebert-uncased")

def get_embedding(text):
    inputs = tokenizer(text, max_length=128, padding="max_length", truncation=True, return_tensors="np")

    input_details = interpreter.get_input_details()
    interpreter.set_tensor(input_details[0]['index'], inputs['input_ids'].astype(np.int64))
    interpreter.set_tensor(input_details[1]['index'], inputs['attention_mask'].astype(np.int64))

    interpreter.invoke()

    output_details = interpreter.get_output_details()
    return interpreter.get_tensor(output_details[0]['index'])[0]

emb = get_embedding("Hello world")
print(f"Embedding shape: {emb.shape}")

Converted by Bombek1 using litert-torch

Bombek1/mobilebert-uncased-litert

google/mobilebert-uncased - LiteRT Optimized

Model Details

Usage

Bombek1/
mobilebert-uncased-litert