intel-optimized-model-for-embeddings-v1

This is a text embedding model model: It maps sentences & paragraphs to a 512 dimensional dense vector space and can be used for tasks like clustering or semantic search. For sample code that uses this model in a torch serve container see Intel-Optimized-Container-for-Embeddings.

Usage

Install the required packages:

pip install -U torch==2.3.1+cpu --extra-index-url https://download.pytorch.org/whl/cpu
pip install -U transformers==4.42.4 intel-extension-for-pytorch==2.3.100

Use the following example below to load the model with the transformers library, tokenize the text, run the model, and apply pooling to the output.

import torch
from transformers import AutoTokenizer, AutoModel
import intel_extension_for_pytorch as ipex

def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0]
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded,
                        1) / torch.clamp(input_mask_expanded.sum(1),
                                        min=1e-9)

# load model
tokenizer = AutoTokenizer.from_pretrained('Intel/intel-optimized-model-for-embeddings-v1')
model = AutoModel.from_pretrained('Intel/intel-optimized-model-for-embeddings-v1', 
                                   torchscript=True)
model.eval()

# do IPEX optimization
batch_size = 1
seq_length=512
vocab_size = model.config.vocab_size
sample_input = {"input_ids": torch.randint(vocab_size, size=[batch_size, seq_length]),
                "token_type_ids": torch.zeros(size=[batch_size, seq_length],
                                                dtype=torch.int),
                "attention_mask": torch.randint(1, size=[batch_size, seq_length])}
text = "This is a test."
model = ipex.optimize(model, level="O1",auto_kernel_selection=True,
                            conv_bn_folding=False, dtype=torch.bfloat16)

with torch.no_grad(), torch.cpu.amp.autocast(cache_enabled=False,
                                                dtype=torch.bfloat16):
    # Compile model
    model = torch.jit.trace(model, example_kwarg_inputs=sample_input,
                                    check_trace=False, strict=False)
    model = torch.jit.freeze(model)
    
    # Call model
    tokenized_text = tokenizer(text, padding=True, truncation=True, return_tensors='pt')
    model_output = model(**tokenized_text)
    sentence_embeddings = mean_pooling(model_output,tokenized_text['attention_mask'])
    embeddings = sentence_embeddings[0].tolist()

# Embeddings output
print(embeddings)

Model Details

Model Description

This model was fine-tuned using the sentence-transformers library based on the BERT-Medium_L-8_H-512_A-8 model using UAE-Large-V1 as a teacher.

Training Datasets