Bombek1/multi-qa-MiniLM-L6-cos-v1-litert

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

# Load model and tokenizer
interpreter = Interpreter(model_path="sentence-transformers_multi-qa-MiniLM-L6-cos-v1.tflite")
interpreter.allocate_tensors()
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/multi-qa-MiniLM-L6-cos-v1")

def get_embedding(text: str) -> np.ndarray:
    """Get sentence embedding for input text."""
    encoded = tokenizer(
        text,
        padding="max_length",
        max_length=512,
        truncation=True,
        return_tensors="np"
    )

    interpreter.set_tensor(input_details[0]["index"], encoded["input_ids"].astype(np.int64))
    interpreter.set_tensor(input_details[1]["index"], encoded["attention_mask"].astype(np.int64))
    interpreter.invoke()

    return interpreter.get_tensor(output_details[0]["index"])[0]

# Example
embedding = get_embedding("Hello, world!")
print(f"Embedding shape: {embedding.shape}")  # (384,)

@inproceedings{reimers-2019-sentence-bert,
    title={Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks},
    author={Reimers, Nils and Gurevych, Iryna},
    booktitle={EMNLP 2019},
    year={2019}
}