How can I speed up the setfit model

🤗Transformers
1

Hello everyone,

I am a newbie to the SETFIT model proposed by Tunstall et al. (2022). (For reference GitHub - huggingface/setfit: Efficient few-shot learning with Sentence Transformers)

I am using the pretrained model named all-mpnet-base-v2.

The training dataset consists of approximately 3600 samples, while the validation dataset comprises around 900 samples. However, with the current code implementation, the entire process is estimated to take over 700 hours. So, I am wondering if anyone can help me reduce the running time. Any help would be greatly appreciated.

I am using pytorch1.10.1 and am not able to upgrade it.

I tried to use “compile”, but it might require pytorch >=2.1.0
model = model.compile(model)

My code:

model = SetFitModel.from_pretrained(
    "sentence-transformers/all-mpnet-base-v2").to(device='cuda')
args = TrainingArguments(
    batch_size=15,
    l2_weight=0.01,
    warmup_proportion = 0.05,
    metric_for_best_model='f1',
    num_epochs=12,
    evaluation_strategy="steps",
    save_strategy="steps",
    # logging_steps=200,
    save_steps=10000,
    eval_steps = 10000,
load_best_model_at_end=True
)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=train,
    eval_dataset=val,
)
trainer.train()