Custom Models do not receive prompt_type or task_name

[ArthurSBianchessi]

Specifically SentenceTransformerWrapper (from mteb/models/sentence_transformer_wrapper.py) does not forward these variables to Custom Model encode.

[Samoed]

Can you provide code to check?

[ArthurSBianchessi]

Sure.
The code in the README.md raises error due to missing task_name

from mteb.encoder_interface import PromptType
import mteb
from mteb import MTEB
import numpy as np

class CustomModel:
    def encode(
        self,
        sentences: list[str],
        task_name: str,
        prompt_type: PromptType | None = None,
        **kwargs,
    ) -> np.ndarray:
        """Encodes the given sentences using the encoder.
        
        Args:
            sentences: The sentences to encode.
            task_name: The name of the task.
            prompt_type: The prompt type to use.
            **kwargs: Additional arguments to pass to the encoder.
            
        Returns:
            The encoded sentences.
        """
        pass

model = CustomModel()
tasks = [mteb.get_task("Banking77Classification") ]            # README doesnt turn the tasks into a list, else it raises an error if get_task is used 
evaluation = MTEB(tasks=tasks)        
evaluation.run(model)

It raises:

TypeError: CustomModel.encode() missing 1 required positional argument: 'task_name'

Printing the inputs of encode with:

class CustomModel:
    def encode(self, sentences, *inputs, **kwargs):
        print(f'inputs: {inputs}')
        print(f'kwargs: {kwargs}')

outputs:

inputs: ()
kwargs: {'prompt_name': None, 'batch_size': 32}

There might be some other variable that is not being forwarded.

[Samoed]

This is correct because task_name is a positional argument and is passed correctly to CustomModel. It would be better to update the example in the README. The CustomModel should function similarly to SentenceTransformerWrapper. You can find more implementations of custom models in the models directory.

[ArthurSBianchessi]

Right now neither task_name nor prompt_type is passed to CustomModels, I did create a pull request, but sent task name as a keyword argument because any encode method that expects taskname as a positional argument does not work, and any other method that currently does not break in mteb amd does not receive *args will break (I expect it to be most models). In my understanding, CustomModel.encode should receive task_name either as a positional or as a keyword argument, but currently it does not receive either.
(In my code it prints (), since no other positional argument are passed. Removing *inputs does not break the code, which would happen if task_name was passed as a positional argument)

This is what I understand that was the expected behavior, sorry if I am misunderstanding something.

[Samoed]

You need to inherit custom model from Wrapper class like:

import mteb
from mteb.encoder_interface import PromptType
import numpy as np
from mteb.models.wrapper import Wrapper

class CustomModel(Wrapper):
    def encode(
        self,
        sentences: list[str],
        task_name: str,
        prompt_type: PromptType | None = None,
        **kwargs,
    ) -> np.ndarray:
        ...

I'll update Readme to make it more relevant. Thank you for report!

[ArthurSBianchessi]

Ok, Thanks!