Using SetFit for regression tasks?

[ericlinML]

I was curious about using SetFit for ordinal Likert scale outcomes (ie IMDB movie reviews). It doesn't seem like an obvious option in the SetFit API. Has anyone tried using SetFit for regression tasks?

[nbertagnolli]

Can you provide some more details about what you're trying to do I might be able to help : ). If you have ordinal data you could treat it like a multiclass problem and setfit should 🤞 just work for that problem. I haven't tried this but taking a look it appears you can set the head of the setfit model to be whatever you want so you could do something like:

from sentence_transformers.losses import CosineSimilarityLoss

from setfit import SetFitModel, SetFitTrainer
from sklearn.linear_model import LinearRegression
from sentence_transformers import SentenceTransformer

# Choose a sentence transformer for your embedding
model_body = SentenceTransformer("sentence-transformers/paraphrase-mpnet-base-v2")

# Choose whatever head you want for regression
model_head = LinearRegression()

# Load a SetFit model from Hub
model = SetFitModel(model_body=model_body, model_head=model_head)

# 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, # The number of text pairs to generate for contrastive learning
    num_epochs=1, # The number of epochs to use for constrastive learning
    column_mapping={"sentence": "text", "label": "label"} # Map dataset columns to text/label expected by trainer
)

# Train and evaluate
trainer.train()

I haven't tried it but I suspect instead of using the default LogisticRegression model head you could replace it with any other reasonable estimator from sklearn. Let me know if you want additional help I'd be happy to take a look if you provide some example input and output data you're struggling with.

[nsorros]

The ability to add a different head would be a nice addition to the documentation actually

[nbertagnolli]

I can open a PR if folks want with examples.

[ericlinML]

My work (sorry can't share data) involves looking at mental health patient audio journals (vastly longer than 512 tokens) that have some self report emotional rating scale attached to them "How sad were you today (rate from 1-5 with 1 being not at all to 5 being extremely sad)?" Generally stuff like that - our biggest hurdle is that we have long texts/transcripts and also small sample size, so I've been trying to think about how we can use something like SetFit in our lab.

My apologies for the late reply!

[kgourgou]

I haven't tried it but I suspect instead of using the default LogisticRegression model head you could replace it with any other reasonable estimator from sklearn. Let me know if you want additional help I'd be happy to take a look if you provide some example input and output data you're struggling with.

I can confirm that this is possible. Since the body is basically trained separately from the head, you can train as usual, then embed your training set with the body model and use the embedded data to train a new head. I've since been using a random forest, for example.

Basically:

setfit_model = SetFitModel(model_body = SentenceTransformer(model_id), model_head=RandomForestClassifier()) 

You should be able to attach any sklearn head as long as it supports the required methods.

[savasy]

But @nbertagnolli SetFit model, associates /pair the instances sharing same label (0 or 1) to augment the set size, which underlies contrastive learning. For the regression problem, we pass numeric values instead of 0/1 labels. How do you expect the models associates/pairs the instances!