Allow setting max length

[blakechi]

This PR is opened for resolving #172 .

• Allow users to set max_length when calling SetFitTrainer.train.
• max_length will be set to the maximum number of tokens the model body can handle if it's too large.
• Add relevant tests.

A snippet of usage:

from datasets import load_dataset
from sentence_transformers.losses import CosineSimilarityLoss

from setfit import SetFitModel, SetFitTrainer


# Load a dataset from the Hugging Face Hub
dataset = load_dataset("sst2")

# Simulate the few-shot regime by sampling 8 examples per class
num_classes = 2
train_dataset = dataset["train"].shuffle(seed=42).select(range(8 * num_classes))
eval_dataset = dataset["validation"]

# Load a SetFit model from Hub
model = SetFitModel.from_pretrained(
    "sentence-transformers/paraphrase-mpnet-base-v2",
    use_differentiable_head=True,
    head_params={"out_features": num_classes},
)

# Create trainer
trainer = SetFitTrainer(
    model=model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss_class=CosineSimilarityLoss,
    metric="accuracy",
    batch_size=16,
    num_iterations=20,
    num_epochs=1,
    column_mapping={"sentence": "text", "label": "label"}
)

# Train the differentiable head only
trainer.unfreeze(keep_body_frozen=True)

# Train with custom `max_length`
trainer.train(
    num_epochs=1,
    batch_size=16,
    body_learning_rate=1e-5,
    learning_rate=1e-2,
    l2_weight=0.0,
    max_length=64,  # set your preferred length here
)

[lewtun]

Thanks for this nice quality of life improvement @blakechi 🔥 !

The PR is looking great, and I've left a few small comments that would be nice to address before we merge

[lewtun]

Should we think about a way to actually test the large value has been overwritten with the model max length?

[blakechi]

you're right. I will modify the test

[blakechi]

just updated it. I tested whether the warning raises up correctly and checked the overwritten value in test_modeling.py.

[lewtun]

Same comment here - should we be testing the behaviour explicitly?

[blakechi]

agree. Will push an update later

[blakechi]

just updated it!

[lewtun]

Sorry for the delay and thanks for iterating @blakechi 🤘

I've left one nit and then think we can merge!

[lewtun]

nice!

[blakechi]

Sorry for the delay and thanks for iterating @blakechi 🤘

I've left one nit and then think we can merge!

No worries at all! Thanks for the review 🙏🏻

[PhilipMay]

After "If" there is a space missing.

            If`None`, will use the maximum number of tokens the model body can accept.

[blakechi]

Thanks for finding that, @PhilipMay!

[lewtun]

Thanks for iterating - looks great!