code: def test_embed_with_proba(self):
      classifier_kwargs = {
          "fine_tuning_arguments": FineTuningArguments(0.2, 0.95),
          "num_epochs": 1,
      }
      clf_factory = TransformerBasedClassificationFactory(
          TransformerModelArguments(
              "sshleifer/tiny-distilroberta-base"
          ),
          self.num_classes,
          kwargs=classifier_kwargs,
      )
  
      train_set = twenty_news_transformers(
          20, num_labels=self.num_classes
      )
  
      clf = clf_factory.new()
      clf.fit(train_set)
  
      embeddings, proba = clf.embed(
          train_set,
          return_proba=True,
          embedding_method=self.embedding_method,
      )
      self.assertEqual(2, len(embeddings.shape))
      self.assertEqual(len(train_set), embeddings.shape[0])
      self.assertEqual(
          clf.model.config.hidden_size, embeddings.shape[1]
      )
      self.assertEqual(len(train_set), proba.shape[0])
  
imports: import unittest
  import pytest
  import warnings
  import numpy
  from unittest import mock
  from unittest.mock import patch, Mock
  from parameterized import parameterized_class
  from scipy.sparse import issparse
  from small_text.integrations.pytorch.exceptions import PytorchNotFoundError
  from small_text.training.early_stopping import EarlyStopping, NoopEarlyStopping
  from small_text.training.metrics import Metric
  from small_text.training.model_selection import ModelSelection, NoopModelSelection
  from tests.utils.datasets import twenty_news_transformers
  from tests.utils.testing import assert_array_not_equal
label_desc: Use image node_id 4 for calling the EmbeddingTest obj's underlying member method code with example usage: obj.test_embed_with_proba() without return types