code: def test_embed_with_proba(self):
      device = "cuda:0"
      classification_kwargs = {
          "use_differentiable_head": self.use_differentiable_head,
          "multi_label": self.multi_label,
          "device": device,
      }
      setfit_model_args = SetFitModelArguments(
          "sentence-transformers/paraphrase-MiniLM-L3-v2"
      )
      clf_factory = SetFitClassificationFactory(
          setfit_model_args,
          self.num_classes,
          classification_kwargs=classification_kwargs,
      )
  
      train_set = twenty_news_text(
          20, num_classes=self.num_classes, multi_label=self.multi_label
      )
  
      clf = clf_factory.new()
      clf.fit(train_set)
  
      with patch.object(
          clf.model.model_body,
          "encode",
          wraps=clf.model.model_body.encode,
      ) as encode_spy:
          embeddings, proba = clf.embed(train_set, return_proba=True)
  
          self.assertEqual(1, encode_spy.call_count)
          self.assertEqual(1, len(encode_spy.call_args_list[0].args))
          self.assertEqual(
              len(train_set.x),
              len(encode_spy.call_args_list[0].args[0]),
          )
          self.assertEqual(1, len(encode_spy.call_args_list[0].kwargs))
          self.assertEqual(
              device, encode_spy.call_args_list[0].kwargs["device"]
          )
  
          self.assertEqual(2, len(embeddings.shape))
          self.assertEqual(len(train_set), embeddings.shape[0])
          # pooling layer is named '1'
          self.assertEqual(
              dict(clf.model.model_body.named_modules())[
                  "1"
              ].pooling_output_dimension,
              embeddings.shape[1],
          )
  
imports: import unittest
  import pytest
  import numpy
  from unittest.mock import patch
  from scipy.sparse import issparse
  from small_text.data.datasets import TextDataset
  from small_text.exceptions import UnsupportedOperationException
  from small_text.integrations.pytorch.exceptions import PytorchNotFoundError
  from sklearn.utils.validation import check_is_fitted
  from tests.utils.datasets import twenty_news_text
label_desc: Use image node_id 3 for calling the _EmbeddingTest obj's underlying member method code with example usage: obj.test_embed_with_proba() without return types