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

[PeterK-end]

I was encountering the same issue and downdgraded torch to 2.3.0. Now I get TypeError: cannot pickle 'torch._C.Generator' object from calling active_learner.save("out/test_model.pkl").

Environment

• Python version: 3.12
• small-text: 2.0.0.dev1
• setfit: 1.1.1
• pytorch: 2.3.0
• transformers: 4.50.0.dev0

Installation (pip, conda, or from source): pip

Script

My script is quite similiar:

from small_text import (
    PoolBasedActiveLearner,
    random_initialization_balanced,
    BreakingTies,
    SubsamplingQueryStrategy
)
from small_text.integrations.transformers.classifiers.setfit import SetFitModelArguments
from small_text.integrations.transformers.classifiers.factories import SetFitClassificationFactory
from setfit import TrainingArguments 
from transformers import AutoModel

# https://huggingface.co/sentence-transformers/paraphrase-mpnet-base-v2 (109M params)
sentence_transformer_model_name = 'sentence-transformers/paraphrase-mpnet-base-v2'
tokenizer = AutoTokenizer.from_pretrained(sentence_transformer_model_name)

setfit_model_args = SetFitModelArguments(sentence_transformer_model_name)
# Factory automates classifier creation for iterative learning
clf_factory = SetFitClassificationFactory(setfit_model_args,
                                          num_classes,
                                          classification_kwargs={
                                              'device': "cuda",
                                              'max_seq_len': 128,
                                              'mini_batch_size': 16,
                                              # https://huggingface.co/docs/setfit/v1.0.1/en/conceptual_guides/sampling_strategies
                                              # https://huggingface.co/docs/setfit/v1.0.1/en/reference/trainer#setfit.TrainingArguments
                                              #'trainer_kwargs': {#'sampling_strategy':'undersampling',
                                              #                    'num_iterations': 1
                                              #                  }
                                              
                                          })
# Define our query strategy
query_strategy = SubsamplingQueryStrategy(BreakingTies())

def initialize_with_coldstart():
    cold_start_labels = np.array(y_expert)

    # Create the labeled dataset for initialization
    labeled_texts = [dataset_hf["train"]['text'][i] for i in initial_indices]
    cold_start_dataset = TextDataset.from_arrays(
        labeled_texts,
        cold_start_labels,
        target_labels=target_labels
    )

    # Full training dataset (labeled + unlabeled)
    training_x = [train_text[i] for i in range(len(dataset_hf["train"])) if i not in initial_indices]
    training_y = np.append(cold_start_labels, np.full(len(training_x), -1))

    # Full training dataset with all samples (some labeled, some -1)
    train = TextDataset.from_arrays(
        labeled_texts + training_x,
        training_y,
        target_labels=target_labels
    )

    # Use the active learner with the labeled dataset only
    setfit_train_kwargs = {'show_progress_bar': True}
    active_learner = PoolBasedActiveLearner(clf_factory, query_strategy, train, fit_kwargs={'setfit_train_kwargs': setfit_train_kwargs})

    active_learner.initialize(np.arange(len(initial_indices)))

    return active_learner, train

active_learner, train = initialize_with_coldstart()

[chschroeder]

@PeterK-end Thank you for the additional info! I just released a new dev version last week, unfortunately this issue was not yet part of it, but it is the very next thing on my list and will be done before 2.0.0.

If this is urgent, you could bypass (or override) the small-text serialization. For the "official" solution I will add save / load methods to all classifiers, where the ones using torch models will delegate to torch.save() / model.load_state_dict() instead of pickling. You could do that as well "from outside".

More information on the underlying issue and for future reference: pytorch/pytorch#43672

[jakobstgl]

Hey @PeterK-end , at some point I also dealt with the TypeError: cannot pickle 'torch._C.Generator' object error. Setting up a fresh environment using torch 2.3.0 solved it for me, though I did not know what exactly caused the problem. However, I could now trace back the error to an old version of the sentence-transformers library. If I use sentence-transformers==3.3.1 it crashes, but with sentence-transformers==3.4.1 it throws no error. Maybe that works for you too until the new saving/loading solutions are implemented.

And thank you again @chschroeder for the torch 2.3.0 workaround!

[PeterK-end]

Thank you for the quick replies! I'm looking forward to the new version. For now @jakobstgl's fix with upgrading sentence-transformers to version 3.4.1 does also work. Thank you!

[chschroeder]

Just to give you an update: this proves to be quite challenging. The best outcome for a quick solution would have been if torch had fixed that bug, but since this has been a problem for a while this might not happen soon.

I now have the following options:

1. Custom serialization provided by small text for each (!) classifier, query strategy, stopping criterion, vector index. Moreover, it needs to be extensible, otherwise it would break on custom implementations. (high effort including future additions, cleanest solution)
2. Workaround: while I have not tested this, maybe we could delete the generator objects that cause this error before serialization and restore them upon deserialization (low effort, ugly).
3. A mix of the above options, where we use the custom serialization for setfit classifiers, but stay with pickle/dill for everything else.

Not sure yet, but with 2.0.0 in mind, where this issue is the largest remaining blocker, I might opt for one of the quick solutions.

[chschroeder]

Looking into this right now:

Issue 1: TypeError: cannot pickle 'ConfigModuleInstance' object

• torch=2.5.1+cu124, datasets=3.2.0, pyarrow=18.1.0:
• thrown when serializing datasets.table.InMemoryTable in Dataset.from_dict()

Issue 2: _pickle.PicklingError: Can't pickle : it's not found as torch._C._dynamo.eval_frame.reset_code

• torch==2.4.1+cu121

Upgrading from setfit 1.1.1 to 1.1.2solved both issues for me.

Solution: Upgrade to setfit 1.1.2.

Since this is overall rather troublesome to fix effectively from small-text, I think the optimal solution here is to increase the minimum requirement for setfit.