fix: add sentence trimming to OpenAIWrapper

[yjoonjang]

#1494
@Samoed

[yjoonjang]

tiktoken library needs to be installed.

[Samoed]

Could you provide the results of the model run to ensure that the parameters were passed correctly?

[yjoonjang]

Could you provide the results of the model run to ensure that the parameters were passed correctly?

Doing some tests right now.

[yjoonjang]

My test code for evaluating Ko-StrategyQA (Korean retrieval task) is:

"""Example script for benchmarking all datasets constituting the MTEB Korean leaderboard & average scores"""
from __future__ import annotations

import os
import logging
from multiprocessing import Process, current_process
import torch

from sentence_transformers import SentenceTransformer
from sentence_transformers.models import StaticEmbedding

import mteb
from mteb import MTEB, get_tasks
from mteb.encoder_interface import PromptType
from mteb.models.sentence_transformer_wrapper import SentenceTransformerWrapper
from mteb.models.instruct_wrapper import instruct_wrapper

import argparse
from dotenv import load_dotenv
# from setproctitle import setproctitle
import traceback

load_dotenv()

logging.basicConfig(level=logging.INFO)

logger = logging.getLogger("main")

TASK_LIST_CLASSIFICATION = []

TASK_LIST_CLUSTERING = []

TASK_LIST_PAIR_CLASSIFICATION = []

TASK_LIST_RERANKING = []

TASK_LIST_RETRIEVAL = [
    "Ko-StrategyQA",
    # "AutoRAGRetrieval",
    # "MIRACLRetrieval",
    # "PublicHealthQA",
    # "BelebeleRetrieval",
    # "MrTidyRetrieval",
    # "MultiLongDocRetrieval",
    # "XPQARetrieval"
]

TASK_LIST_STS = []

TASK_LIST = (
    TASK_LIST_CLASSIFICATION
    + TASK_LIST_CLUSTERING
    + TASK_LIST_PAIR_CLASSIFICATION
    + TASK_LIST_RERANKING
    + TASK_LIST_RETRIEVAL
    + TASK_LIST_STS
)

model_names = [
    "openai/text-embedding-3-large", # 8191
]

def evaluate_model(model_name, gpu_id):
    try:
        os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
        
        model = None
        if not os.path.exists(model_name):
            if "m2v" in model_name:
                static_embedding = StaticEmbedding.from_model2vec(model_name)
                model = SentenceTransformer(modules=[static_embedding])
            else:
                model = mteb.get_model(model_name)
        else:
            file_name = os.path.join(model_name, "model.safetensors")
            if os.path.exists(file_name):
                if "m2v" in model_name:
                    static_embedding = StaticEmbedding.from_model2vec(model_name)
                    model = SentenceTransformer(modules=[static_embedding])
                else:
                    model = mteb.get_model(model_name)

        if model:
            # setproctitle(f"{model_name}-{gpu_id}")
            print(f"Running task: {TASK_LIST} / {model_name} on GPU {gpu_id} in process {current_process().name}")
            evaluation = MTEB(
                tasks=get_tasks(tasks=TASK_LIST, languages=["kor-Kore", "kor-Hang", "kor_Hang"])
            )
            # 48GB VRAM 기준 적합한 batch sizes
            if "multilingual-e5" in model_name:
                batch_size = 256
            elif "jina" in model_name:
                batch_size = 8
            elif "bge-m3" in model_name:
                batch_size = 32
            elif "gemma2" in model_name:
                batch_size = 256 
            else:
                batch_size = 64

            if args.quantize: # quantized model의 경우
                evaluation.run(
                    model,
                    output_folder=f"/data_x/EMBEDDING/RESULTS/{model_name}-quantized",
                    encode_kwargs={"batch_size": batch_size, "precision": "binary"},
                )
            else:
                evaluation.run(
                    model,
                    output_folder=f"results/{model_name}",
                    encode_kwargs={"batch_size": batch_size},
                )
    except Exception as ex:
        print(ex)
        traceback.print_exc()

if __name__ == "__main__":
    gpu_id = 0
    evaluate_model(model_names[0], gpu_id)

and got these results:

{
        "precision_at_1": 0.68243,
        "precision_at_10": 0.15,
        "precision_at_100": 0.01711,
        "ndcg_at_1": 0.68243,
        "ndcg_at_10": 0.73607,
        "ndcg_at_100": 0.76238,
        "recall_at_1": 0.44298,
        "recall_at_10": 0.80961,
        "recall_at_100": 0.91126,
}

(Others excluded)

[yjoonjang]

After deleting changes for ModelMetadata, I tested on AutoRAGRetrieval, which is also a Korean retrieval task.
The results are:

{
        "precision_at_1": 0.58772,
        "precision_at_10": 0.09386,
        "precision_at_100": 0.00991,
        "ndcg_at_1": 0.58772,
        "ndcg_at_10": 0.76549,
        "ndcg_at_100": 0.77714,
        "recall_at_1": 0.58772,
        "recall_at_10": 0.9386,
        "recall_at_100": 0.99123,
}

(Others excluded)

The results shows that the parameters are passed well.

[Samoed]

Great!

[KennethEnevoldsen]

Looks good. A few minor changes.

Might be worth rerunning openai given these new changes (@Muennighoff I imagine that we could use the key here)

[Muennighoff]

I forgot to push this but I actually also implemented this locally. IIUC the things that'd be nice to add here such that I can discard my implementation would be:

• retry loop (cleaner than mine hopefully)
• only require tiktoken for specific embedding models
• request packages in init not in encode
• make batch size configurable

class OpenAIWrapper(Wrapper):
    def __init__(self, model_name: str, embed_dim: int | None = None, **kwargs) -> None:
        requires_package(self, "openai", "Openai text embedding")
        from openai import OpenAI
        if model_name in ["text-embedding-3-small", "text-embedding-3-large"]:
            requires_package(self, "tiktoken", "Openai text embedding")
            import tiktoken
            self.max_tokens = 8191
            self.encoding = tiktoken.get_encoding("cl100k_base")

        self._client = OpenAI()
        self._model_name = model_name
        self._embed_dim = embed_dim

    def truncate_text_tokens(self, text):
        """Truncate a string to have `max_tokens` according to the given encoding."""
        return self.encoding.encode(text)[:self.max_tokens]

    def encode(self, sentences: list[str], **kwargs: Any) -> np.ndarray:
        from openai import NotGiven

        if self._model_name == "text-embedding-ada-002" and self._embed_dim is not None:
            logger.warning(
                "Reducing embedding size available only for text-embedding-3-* models"
            )

        max_batch_size = kwargs.get("batch_size", 256)
        sublists = [
            sentences[i : i + max_batch_size]
            for i in range(0, len(sentences), max_batch_size)
        ]

        all_embeddings = []

        for sublist in sublists:
            if self._model_name in ["text-embedding-3-small", "text-embedding-3-large"]:
                sublist = [self.truncate_text_tokens(s) for s in sublist]
            try:
                # Takes either strings or tokenized inputs
                response = self._client.embeddings.create(
                    input=sublist,
                    model=self._model_name,
                    encoding_format="float",
                    dimensions=self._embed_dim or NotGiven(),
                )
            except Exception as e:
                # import pdb; pdb.set_trace()
                # Sleep
                print("Sleeping for 10 seconds due to error", e)
                import time
                time.sleep(10)
                try:
                    response = self._client.embeddings.create(
                        input=sublist,
                        model=self._model_name,
                        encoding_format="float",
                        dimensions=self._embed_dim or NotGiven(),
                    )
                except Exception as e:
                    print("Sleeping for 60 seconds due to error", e)
                    time.sleep(60)
                    response = self._client.embeddings.create(
                        input=sublist,
                        model=self._model_name,
                        encoding_format="float",
                        dimensions=self._embed_dim or NotGiven(),
                    )
            all_embeddings.extend(self._to_numpy(response))

        return np.array(all_embeddings)

    def _to_numpy(self, embedding_response) -> np.ndarray:
        return np.array([e.embedding for e in embedding_response.data])

[Muennighoff]

@KennethEnevoldsen I don't think we need to rerun things as I already ran them with the truncation per my implementation above; Makes more sense to use for new tasks I think. Another thing we could run is using the Matryoshka parameter to change the embedding size and benchmark the large model but with a smaller emb size

[yjoonjang]

I forgot to push this but I actually also implemented this locally. IIUC the things that'd be nice to add here such that I can discard my implementation would be:

• retry loop (cleaner than mine hopefully)
• only require tiktoken for specific embedding models
• request packages in init not in encode
• make batch size configurable

class OpenAIWrapper(Wrapper):
    def __init__(self, model_name: str, embed_dim: int | None = None, **kwargs) -> None:
        requires_package(self, "openai", "Openai text embedding")
        from openai import OpenAI
        if model_name in ["text-embedding-3-small", "text-embedding-3-large"]:
            requires_package(self, "tiktoken", "Openai text embedding")
            import tiktoken
            self.max_tokens = 8191
            self.encoding = tiktoken.get_encoding("cl100k_base")

        self._client = OpenAI()
        self._model_name = model_name
        self._embed_dim = embed_dim

    def truncate_text_tokens(self, text):
        """Truncate a string to have `max_tokens` according to the given encoding."""
        return self.encoding.encode(text)[:self.max_tokens]

    def encode(self, sentences: list[str], **kwargs: Any) -> np.ndarray:
        from openai import NotGiven

        if self._model_name == "text-embedding-ada-002" and self._embed_dim is not None:
            logger.warning(
                "Reducing embedding size available only for text-embedding-3-* models"
            )

        max_batch_size = kwargs.get("batch_size", 256)
        sublists = [
            sentences[i : i + max_batch_size]
            for i in range(0, len(sentences), max_batch_size)
        ]

        all_embeddings = []

        for sublist in sublists:
            if self._model_name in ["text-embedding-3-small", "text-embedding-3-large"]:
                sublist = [self.truncate_text_tokens(s) for s in sublist]
            try:
                # Takes either strings or tokenized inputs
                response = self._client.embeddings.create(
                    input=sublist,
                    model=self._model_name,
                    encoding_format="float",
                    dimensions=self._embed_dim or NotGiven(),
                )
            except Exception as e:
                # import pdb; pdb.set_trace()
                # Sleep
                print("Sleeping for 10 seconds due to error", e)
                import time
                time.sleep(10)
                try:
                    response = self._client.embeddings.create(
                        input=sublist,
                        model=self._model_name,
                        encoding_format="float",
                        dimensions=self._embed_dim or NotGiven(),
                    )
                except Exception as e:
                    print("Sleeping for 60 seconds due to error", e)
                    time.sleep(60)
                    response = self._client.embeddings.create(
                        input=sublist,
                        model=self._model_name,
                        encoding_format="float",
                        dimensions=self._embed_dim or NotGiven(),
                    )
            all_embeddings.extend(self._to_numpy(response))

        return np.array(all_embeddings)

    def _to_numpy(self, embedding_response) -> np.ndarray:
        return np.array([e.embedding for e in embedding_response.data])

Woah, sorry I saw this comment too late. I think this code is better because I’ve also experienced an error related to too many requests. Adding a sleep might help.

Can I implement this in my code, or could you revise the code I wrote? Let me know if there are any issues with authorization. @Muennighoff

[Muennighoff]

I forgot to push this but I actually also implemented this locally. IIUC the things that'd be nice to add here such that I can discard my implementation would be:

• retry loop (cleaner than mine hopefully)
• only require tiktoken for specific embedding models
• request packages in init not in encode
• make batch size configurable

class OpenAIWrapper(Wrapper):
    def __init__(self, model_name: str, embed_dim: int | None = None, **kwargs) -> None:
        requires_package(self, "openai", "Openai text embedding")
        from openai import OpenAI
        if model_name in ["text-embedding-3-small", "text-embedding-3-large"]:
            requires_package(self, "tiktoken", "Openai text embedding")
            import tiktoken
            self.max_tokens = 8191
            self.encoding = tiktoken.get_encoding("cl100k_base")

        self._client = OpenAI()
        self._model_name = model_name
        self._embed_dim = embed_dim

    def truncate_text_tokens(self, text):
        """Truncate a string to have `max_tokens` according to the given encoding."""
        return self.encoding.encode(text)[:self.max_tokens]

    def encode(self, sentences: list[str], **kwargs: Any) -> np.ndarray:
        from openai import NotGiven

        if self._model_name == "text-embedding-ada-002" and self._embed_dim is not None:
            logger.warning(
                "Reducing embedding size available only for text-embedding-3-* models"
            )

        max_batch_size = kwargs.get("batch_size", 256)
        sublists = [
            sentences[i : i + max_batch_size]
            for i in range(0, len(sentences), max_batch_size)
        ]

        all_embeddings = []

        for sublist in sublists:
            if self._model_name in ["text-embedding-3-small", "text-embedding-3-large"]:
                sublist = [self.truncate_text_tokens(s) for s in sublist]
            try:
                # Takes either strings or tokenized inputs
                response = self._client.embeddings.create(
                    input=sublist,
                    model=self._model_name,
                    encoding_format="float",
                    dimensions=self._embed_dim or NotGiven(),
                )
            except Exception as e:
                # import pdb; pdb.set_trace()
                # Sleep
                print("Sleeping for 10 seconds due to error", e)
                import time
                time.sleep(10)
                try:
                    response = self._client.embeddings.create(
                        input=sublist,
                        model=self._model_name,
                        encoding_format="float",
                        dimensions=self._embed_dim or NotGiven(),
                    )
                except Exception as e:
                    print("Sleeping for 60 seconds due to error", e)
                    time.sleep(60)
                    response = self._client.embeddings.create(
                        input=sublist,
                        model=self._model_name,
                        encoding_format="float",
                        dimensions=self._embed_dim or NotGiven(),
                    )
            all_embeddings.extend(self._to_numpy(response))

        return np.array(all_embeddings)

    def _to_numpy(self, embedding_response) -> np.ndarray:
        return np.array([e.embedding for e in embedding_response.data])

Woah, sorry I saw this comment too late. I think this code is better because I’ve also experienced an error related to too many requests. Adding a sleep might help.

Can I implement this in my code, or could you revise the code I wrote? Let me know if there are any issues with authorization. @Muennighoff

Ofc feel free to implement it in your code and take the bits from my code that are useful to you; then after you commit it hopefully we can merge this PR! 🚀

[yjoonjang]

I changed my codes thanks to great reviews and comments.
For AutoRAGRetrieval (Which is Korean retrieval task),
text-embedding-3-large scored

"ndcg_at_1": 0.57018,
"ndcg_at_3": 0.71157,
"ndcg_at_5": 0.74102,
"ndcg_at_10": 0.76131,

and text-embedding-ada-002 scored

"ndcg_at_1": 0.47368,
"ndcg_at_3": 0.61831,
"ndcg_at_5": 0.63605,
"ndcg_at_10": 0.6677,

[KennethEnevoldsen]

I think everything looks good here. Will merge it in - thanks for the contribution!

[Muennighoff]

If OpenAI embedding credit for MTEB-related benchmarking would be useful to anyone here in this thread, let me know! We have 1K+ USD in credits expiring Jan 31st :)