# %pip install setfit[optuna] matplotlibimport numpy as np
from datasets import load_dataset
from sentence_transformers.losses import CosineSimilarityLoss
from setfit import SetFitModel, SetFitTrainerdef encode_labels(record):
    return {"labels": [record[feature] for feature in features]}


num_samples = 8
dataset = load_dataset("ethos", "multilabel")

features = dataset["train"].column_names
features.remove("text")
samples = np.concatenate([np.random.choice(np.where(dataset["train"][f])[0], num_samples) for f in features])
dataset = dataset.map(encode_labels)
train_dataset = dataset["train"].select(samples)
eval_dataset = dataset["train"].select(np.setdiff1d(np.arange(len(dataset["train"])), samples))model_id = "sentence-transformers/paraphrase-mpnet-base-v2"


def make_model(params=None):
    multi_target_strategy = params["multi_target_strategy"] if params else "one-vs-rest"
    return SetFitModel.from_pretrained(
        model_id, multi_target_strategy=multi_target_strategy
    )trainer = SetFitTrainer(
    model_init=make_model,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    loss_class=CosineSimilarityLoss,
    num_epochs=1,
    num_iterations=5,
    column_mapping={"text": "text", "labels": "label"},
)def hyperparameter_search_function(trial):
    return {
        "learning_rate": trial.suggest_float("learning_rate", 1e-5, 1e-3, log=True),
        "batch_size": trial.suggest_categorical("batch_size", [4, 8, 16, 32]),
        "multi_target_strategy": trial.suggest_categorical("multi_target_strategy", ["one-vs-rest", "multi-output", "classifier-chain"])
    }best = trainer.hyperparameter_search(hyperparameter_search_function, n_trials=10)
bestfrom optuna.visualization.matplotlib import plot_param_importances

plot_param_importances(best.backend);trainer.apply_hyperparameters(best.hyperparameters, final_model=True) # replaces model_init with a fixed model
trainer.train()metrics = trainer.evaluate()
best.objective, metrics