setfit model performance advice

Hello gents,

I was hoping I can get a second opinion about a situation I am facing while using setfit for a multi class classification use case.
The dataset is small with 255 samples across 9 classes. It suffers from both ambiguity (overlapping) and class imbalance. I partially remedied the overlaps by merging classes together. The imbalance is still present but the new dataset version is such as , except for one class, all the other have at least 10 samples per class as shown here below:

I am performing a stratified split training 80%, validation 10%, testing 10%. The train, validation and test splits are preserving the initial class imbalance observed in the dataset post-merge.

My model is an embeddings model with a logistic regression head, constructed using this code:

```python
model_body = SentenceTransformer("BAAI/bge-base-en-v1.5")

# Choose Logistic Regression as the classification head
model_head = LogisticRegression(class_weight="balanced")

# Create a SetFit model, combining the feature extractor and classification head
model = SetFitModel(model_body, model_head)
model.labels = categories
# labels = list(set(dataset['train']['label']))

# Load a SetFit model from Hub
model: SetFitModel = SetFitModel.from_pretrained(
"BAAI/bge-small-en-v1.5",
)
```

I am then using the following code (with the latest hyper parameters choice) to train the setfit model:

```python
# Create Training Arguments
args = TrainingArguments(
# When an argument is a tuple, the first value is for training the embeddings,
# and the latter is for training the differentiable classification head:
batch_size=(32, 2),
num_iterations=10,
num_epochs=(5, 16),
body_learning_rate=(1e-5, 1e-5),
head_learning_rate=2e-2,
end_to_end=True,
show_progress_bar=False,
report_to="none",
logging_strategy="steps",
logging_steps=50,
eval_steps=50,
output_dir= f"{current_path}/checkpoints",
logging_dir = f"{current_path}/runs",
seed=40
)

# Create Trainer
trainer = Trainer(
model=model,
args=args,
train_dataset=dataset['train'],
eval_dataset=dataset['validation'],
metric="accuracy",
)

# Train and evaluate
trainer.train()
```

This is the logging output I get during training:

the evaluation step results in an accuracy score of 0.6153846153846154

I am plotting the embedding and eval embedding loss curves:


when I predict against the test split , I get the following results:

Looking at these results but especially the embedding and eval embedding curves, it's obvious the training routine needs improvement. I was initially suffering from a eval embedding loss curve that was flat, but the problem was in the splitting and that problem is now solved. But the curves shouldn't look like this.
I also suspect the classifier might not be learning enough (it's just a suspicion based on digging at the predictions it made against the test split).

I welcome any helpful suggestions.

PS:

1) I have already tried suggested hyperparameters, I read on other GitHub issues here like:
batch_size=(16, 2)
num_epochs=(3,16),
body_learning_rate=(2e-5, 1e-5),
head_learning_rate=1e-2,

but I got this:

and worse evaluation and test prediction scores

2) already tried optuna but do not seem to get anywhere.