Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Passing arbitrary keyword arguments to submodules #2539

Open
wants to merge 11 commits into
base: master
Choose a base branch
from

Conversation

Markus28
Copy link

@Markus28 Markus28 commented Mar 12, 2024

This PR implements a prototype of the proposal #2538. I.e., it allows arbitrary keyword arguments to be passed through SentenceTransformer.encode to the model.

Users may now specify in modules.json what keyword arguments are expected by which module. The SentenceTransformer class no longer inherits from nn.Sequential, but instead from nn.ModuleDict. We implement the forward method ourselves and distribute keyword arguments to the modules.

A model may then provide a modules.json file like this:

[
  {
    "idx": 0,
    "name": "0",
    "path": "",
    "type": "sentence_transformers.models.Transformer",
    "kwargs": ["task", "embedding_dim", "foobar"]
  },
  {
    "idx": 1,
    "name": "1",
    "path": "1_Pooling",
    "type": "sentence_transformers.models.Pooling"
  }
]

and users may use the sentence-transformer model like this:

model = SentenceTransformer('<MODEL>', trust_remote_code=True)
model.encode(['Hello world'], task='sts', embedding_dim=32, foobar=0)

Todo

  • These changes will be breaking if there is someone using the .append method, which was previously inherited from nn.Sequential. However, we could also implement it here
  • It would make sense to also implement arbitrary kwargs for the tokenizer

@bwanglzu
Copy link
Contributor

@tomaarsen please let us know how do you think about this PR :) thanks

Bo

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

2 participants