Active learner save method does not work for SetFit

[jakobstgl]

Bug description

I set up a Classification Factory for a SetFit Model, initialized the PoolBasedActiveLearner and started the Query Process successfully. However, when trying to save the Active Learner to disk for later usage, I ran into TypeError: cannot pickle 'ConfigModuleInstance' object

Steps to reproduce

def initialize_active_learner(active_learner, y_train):
    indices_initial = range(0,50)
    active_learner.initialize(np.array(indices_initial))
    
    return indices_initial


num_classes = 2
sentence_transformer_model_name = 'sentence-transformers/paraphrase-mpnet-base-v2'
setfit_model_args = SetFitModelArguments(sentence_transformer_model_name)
clf_factory = SetFitClassificationFactory(
    setfit_model_args,
    num_classes,
    classification_kwargs = {
    'device': 'cuda',
    'max_seq_len': 64,
    'mini_batch_size': 8
}
)
active_learner = PoolBasedActiveLearner(clf_factory, query_strategy, dataset = dataset)

query_strategy = SubsamplingQueryStrategy(BreakingTies(), subsample_size=10000) 
active_learner = PoolBasedActiveLearner(clf_factory, query_strategy, dataset = dataset)

indices_initial = initialize_active_learner(active_learner, dataset.y)
indices_labeled = indices_initial

active_learner.save("test.pkl")

Expected behavior

Serialize and save the active learner in the same way as it works with a regular transformer-based active learner.

Environment:

Python version: 3.9.6
small-text version: 2.0.0.dev1
small-text integrations (e.g., transformers): setfit 1.1.0, transformers 4.45.2
PyTorch version (if applicable): 2.5.0+cu124

Installation (pip, conda, or from source): pip
CUDA version (if applicable): 12.6

Addition information

[chschroeder]

Hi @jakobstgl, thanks for the detailed report!

I quickly tested this and it seems to depend on the torch version:

• With 2.3.0 everything seems to work.
• Wth 2.4.1 I get: _pickle.PicklingError: Can't pickle <built-in function reset_code>: it's not found as torch._C._dynamo.eval_frame.reset_code.
• With 2.5.1, I can reproduce: TypeError: cannot pickle 'ConfigModuleInstance' object

As a quick solution, you might try to downgrade torch to 2.3.0. For a general solution, I need to investigate what changed in torch.

[chschroeder]

Still looking into this. Best leads that I have are linked here pytorch/pytorch#126154.

The actual error looks like a problem that is out of my control though. Best I can do from the small-text point of view, is to provide a better serialization method, and stop the pickling.

Short-term workaround: torch 2.3.0