jinaai

/

jina-embeddings-v2-base-en

like 681

Ok, I think I got a walk around:

!git clone https://huggingface.co/jinaai/jina-bert-implementation
!mv jina-bert-implementation jina_bert_implementation
!touch jina_bert_implementation/__init__.py

from jina_bert_implementation.modeling_bert import JinaBertModel

checkpoint = "my_checkpoint"
model = JinaBertModel.from_pretrained(checkpoint)
model.to(device)

from transformers import AutoTokenizer
import torch
#Mean Pooling - Take attention mask into account for correct averaging
def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
    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 from HuggingFace Hub
tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-embeddings-v2-base-en')

# Tokenize sentences
encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
encoded_input = {
    key: val.to(device) for key, val in encoded_input.items()
}
# Compute token embeddings
with torch.no_grad():
    model_output = model(**encoded_input)

# Perform pooling. In this case, max pooling.
sentences_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])

Still would appreciate the help, because I want to load it as SentenceTransformer for ease of use.

hi @Maiia can you manually edit the SentenceTransformer class, add trust_remote_code=True when sbert doing the AutoModel.from_pretrained(...) thingy?

i think in SBert main branch they support it, not in the latest pypi release.

Was not able to find where to change it, but I adapted the function and created a class similar to SentenceTransformer (at least it does the encoding efficiently)
Maybe someone else finds it useful:

from tqdm.notebook import trange
import numpy as np
import torch

from transformers import AutoModel, AutoTokenizer

def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0] #First element of model_output contains all token embeddings
    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)

tokenizer = AutoTokenizer.from_pretrained('jinaai/jina-embeddings-v2-base-en')

class JinaSentEmbedder(AutoModel):
    def __init__(self, path):
        
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = AutoModel.from_pretrained(
          path, 
          trust_remote_code=True
        )
          
        self.model = self.model.to(self.device)
        self.tokenize = AutoTokenizer.from_pretrained(
            "jinaai/jina-embeddings-v2-base-en"
        )
        
    def _text_length(self, text):
        if isinstance(text, dict):              #{key: value} case
            return len(next(iter(text.values())))
        elif not hasattr(text, '__len__'):      #Object has no len() method
            return 1
        elif len(text) == 0 or isinstance(text[0], int):    #Empty string or list of ints
            return len(text)
        else:
            return sum([len(t) for t in text])
    def encode(self, sentences,
               batch_size = 32,
               show_progress_bar = None,
               output_value: str = 'sentence_embedding',
               convert_to_numpy: bool = True,
               convert_to_tensor: bool = False,
               device: str = None,
               normalize_embeddings: bool = False):

        self.model.eval()

        if convert_to_tensor:
            convert_to_numpy = False

        if output_value != 'sentence_embedding':
            convert_to_tensor = False
            convert_to_numpy = False

        input_was_string = False
        if isinstance(sentences, str) or not hasattr(sentences, '__len__'): #Cast an individual sentence to a list with length 1
            sentences = [sentences]
            input_was_string = True

        all_embeddings = []
        length_sorted_idx = np.argsort([-self._text_length(sen) for sen in sentences])
        sentences_sorted = [sentences[idx] for idx in length_sorted_idx]

        for start_index in trange(0, len(sentences), batch_size, desc="Batches", disable=not show_progress_bar):
            sentences_batch = sentences_sorted[start_index:start_index+batch_size]
            encoded_input = self.tokenize(sentences_batch, padding=True, truncation=True, return_tensors='pt')
            encoded_input = {key: val.to(self.device) for key, val in encoded_input.items()}

            with torch.no_grad():
                model_output = self.model(**encoded_input)
                sentences_embeddings = mean_pooling(model_output, encoded_input['attention_mask'])
            all_embeddings.extend(sentences_embeddings)

        all_embeddings = [all_embeddings[idx].cpu() for idx in np.argsort(length_sorted_idx)]

        if convert_to_tensor:
            all_embeddings = torch.stack(all_embeddings)
        elif convert_to_numpy:
            all_embeddings = np.asarray([emb.numpy() for emb in all_embeddings])

        if input_was_string:
            all_embeddings = all_embeddings[0]

        return all_embeddings

Hi @Maiia, could you please share the code you utilized for fine-tuning this model?
Thank you in advance!

@metalwhale
Hello, it's just normal SentenceTransformers fine-tuning, I have marked up pairs of phrases with labels (so phrase1, phrase2, label) where "label" can be either "pos" or "neg"

title_df = pl.DataFrame({
    "title 1": [el[0] for el in dedup_negatives] + [el[0] for el in hard_positives],
    'title 2': [el[1] for el in dedup_negatives] + [el[1] for el in hard_positives],
    'label': ['neg'] * len(dedup_negatives) + ['pos'] * len(hard_positives)
})

for _ in range(5):
    title_df = title_df.sample(fraction=1, shuffle=True)
train_df, val_df = train_test_split(title_df, random_state=42,
                                    test_size=0.1,
                                    stratify=title_df['label'].to_list())
train_examples = []
for row in train_df.iter_rows(named=True):
    train_examples.append(
        InputExample(texts=[row['title 1'], row['title 2']],
                     label=torch.tensor(1 if row['label'] == 'pos' else 0).to(torch.float32))
    )
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)
train_loss = losses.CosineSimilarityLoss(model)

sentences1 = val_df['title 1'].to_list()
sentences2 = val_df['title 2'].to_list()
scores = [torch.tensor(1 if el == 'pos' else 0).to(torch.float32) for el in val_df['label'].to_list()]

evaluator = EmbeddingSimilarityEvaluator(sentences1, sentences2, scores)

model.fit(train_objectives=[(train_dataloader, train_loss)], epochs=1,
          warmup_steps=len(train_dataloader)//10,
          evaluator=evaluator, evaluation_steps=len(train_dataloader)//10)

@metalwhale I hope it helps, here is the documentation: https://www.sbert.net/docs/training/overview.html

@Maiia thank you so much for your kind help. I really appreciate it!

Hi @Maiia , I followed your step to add a JinaSentEmbedder class and load the jina model with the following code:

jina_path = './jina-embeddings-v2-base-code'
model = JinaSentEmbedder(jina_path)
...
model.fit(...)

It seems fit is not a part of JinaSentEmbedder, I read the source code of sentence-transformers and found that fit is implemented in SentenceTransformer.
Does this mean I should copy a SentenceTransformer class and patch the functions implemented in JinaSentEmbedder? Or is there another way to load JinaSentEmbedderas SentenceTransformer?
This might be a dumb question as I am new to transformers etc. Thank you in advance!

@shijy16 a lot has happened since this issue was created. You can now load this model and finetune with it with normal Sentence Transformers. A complete training script may look like this:

import logging
from datasets import load_dataset, Dataset
from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
    SentenceTransformerModelCardData,
)
from sentence_transformers.losses import MultipleNegativesRankingLoss
from sentence_transformers.training_args import BatchSamplers
from sentence_transformers.evaluation import InformationRetrievalEvaluator


logging.basicConfig(
    format="%(asctime)s - %(message)s", datefmt="%Y-%m-%d %H:%M:%S", level=logging.INFO
)

# 1. Load a model to finetune with 2. (Optional) model card data
model = SentenceTransformer(
    "jinaai/jina-embeddings-v2-base-en",
    trust_remote_code=True,
    model_card_data=SentenceTransformerModelCardData(
        language="en",
        license="apache-2.0",
        model_name="jina-embeddings-v2-base-en trained on Natural Questions pairs",
    ),
)
model_name = "jina-v2-base-natural-questions"

# 3. Load a dataset to finetune on
dataset = load_dataset("sentence-transformers/natural-questions", split="train")
dataset = dataset.add_column("id", range(len(dataset)))
train_dataset: Dataset = dataset.select(range(90_000))
eval_dataset: Dataset = dataset.select(range(90_000, len(dataset)))

# 4. Define a loss function
loss = MultipleNegativesRankingLoss(model)


# 5. (Optional) Specify training arguments
args = SentenceTransformerTrainingArguments(
    # Required parameter:
    output_dir=f"models/{model_name}",
    # Optional training parameters:
    num_train_epochs=1,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    learning_rate=2e-5,
    warmup_ratio=0.1,
    fp16=False,  # Set to False if you get an error that your GPU can't run on FP16
    bf16=True,  # Set to True if you have a GPU that supports BF16
    batch_sampler=BatchSamplers.NO_DUPLICATES,  # MultipleNegativesRankingLoss benefits from no duplicate samples in a batch
    # Optional tracking/debugging parameters:
    eval_strategy="steps",
    eval_steps=200,
    save_strategy="steps",
    save_steps=200,
    save_total_limit=2,
    logging_steps=200,
    logging_first_step=True,
    run_name=model_name,  # Will be used in W&B if `wandb` is installed
)

# 6. (Optional) Create an evaluator & evaluate the base model
# The full corpus, but only the evaluation queries
queries = dict(zip(eval_dataset["id"], eval_dataset["query"]))
corpus = {cid: dataset[cid]["answer"] for cid in range(10_000)} | {cid: dataset[cid]["answer"] for cid in eval_dataset["id"]}
relevant_docs = {qid: {qid} for qid in eval_dataset["id"]}
dev_evaluator = InformationRetrievalEvaluator(
    corpus=corpus,
    queries=queries,
    relevant_docs=relevant_docs,
    show_progress_bar=True,
    name="natural-questions-dev",
    batch_size=8,
)
dev_evaluator(model)

# 7. Create a trainer & train
trainer = SentenceTransformerTrainer(
    model=model,
    args=args,
    train_dataset=train_dataset.remove_columns("id"),
    eval_dataset=eval_dataset.remove_columns("id"),
    loss=loss,
    evaluator=dev_evaluator,
)
trainer.train()

# (Optional) Evaluate the trained model on the evaluator after training
dev_evaluator(model)

# 8. Save the trained model
model.save_pretrained(f"models/{model_name}/final")

# 9. (Optional) Push it to the Hugging Face Hub
model.push_to_hub(f"{model_name}")

• Tom Aarsen

@tomaarsen This is exactly the answer I am looking for. I really appreciate it. Thank you!

thanks @tomaarsen for the quick reply!