From 4935dd1cb905d6421d53af6e48482dd21fba5a45 Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Tue, 28 Jan 2025 12:35:06 +0000
Subject: [PATCH 01/19] update

---
 README.md                                     |  3 +
 .../quantization_aware_training/Usage.md      | 67 +++++++++++++++----
 2 files changed, 57 insertions(+), 13 deletions(-)

diff --git a/README.md b/README.md
index 4347feec8cf..65bbb7f3e5c 100644
--- a/README.md
+++ b/README.md
@@ -201,6 +201,9 @@ def transform_fn(data_item):
 calibration_dataset = nncf.Dataset(val_dataset, transform_fn)
 # Step 3: Run the quantization pipeline
 quantized_model = nncf.quantize(model, calibration_dataset)
+# Step 4: Remove auxiliary layers and operations added during the quantization process,
+# resulting in a clean, fully quantized model ready for deployment.
+stripped_model = nncf.strip(quantized_model)
 ```
 
 </details>
diff --git a/docs/usage/training_time_compression/quantization_aware_training/Usage.md b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
index 3a5fbffb096..b9f149bee82 100644
--- a/docs/usage/training_time_compression/quantization_aware_training/Usage.md
+++ b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
@@ -1,7 +1,7 @@
-# Use NNCF for Quantization Aware Training in PyTorch
+# Use NNCF for Quantization Aware Training
 
-This is a step-by-step tutorial on how to integrate the NNCF package into the existing PyTorch project (please see the [TensorFlow quantization documentation](../other_algorithms/LegacyQuantization.md) for integration tutorial for the existing TensorFlow project).
-The use case implies that the user already has a training pipeline that reproduces training of the model in the floating  point precision and pretrained model.
+This is a step-by-step tutorial on how to integrate the NNCF package into the existing PyTorch or TensorFlow projects.
+The use case implies that the user already has a training pipeline that reproduces training of the model in the floating point precision and pretrained model.
 The task is to prepare this model for accelerated inference by simulating the compression at train time.
 Please refer to this [document](/docs/usage/training_time_compression/other_algorithms/LegacyQuantization.md) for details of the implementation.
 
@@ -11,11 +11,24 @@ Please refer to this [document](/docs/usage/training_time_compression/other_algo
 
 Quantize the model using the [Post Training Quantization](../../post_training_compression/post_training_quantization/Usage.md) method.
 
+<details open><summary><b>PyTorch</b></summary>
+
 ```python
 model = TorchModel() # instance of torch.nn.Module
 quantized_model = nncf.quantize(model, ...)
 ```
 
+</details>
+
+<details><summary><b>TensorFlow</b></summary>
+
+```python
+model = TensorFlowModel() # instance of tf.keras.Model
+quantized_model = nncf.quantize(model, ...)
+```
+
+</details>
+
 ### Step 2: Run the training pipeline
 
 At this point, the NNCF is fully integrated into your training pipeline.
@@ -27,18 +40,39 @@ Important points you should consider when training your networks with compressio
 
 ### Step 3: Export the compressed model
 
-After the compressed model has been fine-tuned to acceptable accuracy and compression stages, you can export it. There are two ways to export a model:
+After the compressed model has been fine-tuned to acceptable accuracy and compression stages, you can export it.
+
+<details open><summary><b>PyTorch</b></summary>
+
+Trace the model via inference in framework operations.
+
+```python
+# To OpenVINO format
+import openvino as ov
+ov_quantized_model = ov.convert_model(quantized_model.cpu(), example_input=dummy_input)
+```
 
-1. Trace the model via inference in framework operations.
+</details>
 
-    ```python
-    # To OpenVINO format
-    import openvino as ov
-    ov_quantized_model = ov.convert_model(quantized_model.cpu(), example_input=dummy_input)
-    ```
+<details><summary><b>TensorFlow</b></summary>
+
+```python
+# To OpenVINO format
+import openvino as ov
+
+# Removes auxiliary layers and operations added during the quantization process,
+# resulting in a clean, fully quantized model ready for deployment.
+stripped_model = nncf.strip(quantized_model)
+
+ov_quantized_model = ov.convert_model(stripped_model, share_weights=False)
+```
+
+</details>
 
 ## Saving and loading compressed models
 
+<details open><summary><b>PyTorch</b></summary>
+
 The complete information about compression is defined by a compressed model and a NNCF config.
 The model characterizes the weights and topology of the network. The NNCF config - how to restore additional modules intoduced by NNCF.
 The NNCF config can be obtained by `quantized_model.nncf.get_config()` on saving and passed to the
@@ -46,8 +80,6 @@ The NNCF config can be obtained by `quantized_model.nncf.get_config()` on saving
 The quantized model saving allows to load quantized modules to the target model in a new python process and
 requires only example input for the target module, corresponding NNCF config and the quantized model state dict.
 
-### Saving and loading compressed models in PyTorch
-
 ```python
 # save part
 quantized_model = nncf.quantize(model, calibration_dataset)
@@ -70,10 +102,14 @@ quantized_model.load_state_dict(state_dict)
 
 You can save the `compressed_model` object `torch.save` as usual: via `state_dict` and `load_state_dict` methods.
 
+</details>
+
 ## Advanced usage
 
 ### Compression of custom modules
 
+<details open><summary><b>PyTorch</b></summary>
+
 With no target model code modifications, NNCF only supports native PyTorch modules with respect to trainable parameter (weight) compressed, such as `torch.nn.Conv2d`.
 If your model contains a custom, non-PyTorch standard module with trainable weights that should be compressed, you can register it using the `@nncf.register_module` decorator:
 
@@ -91,4 +127,9 @@ If registered module should be ignored by specific algorithms use `ignored_algor
 
 In the example above, the NNCF-compressed models that contain instances of `MyModule` will have the corresponding modules extended with functionality that will allow NNCF to quantize the `weight` parameter of `MyModule` before it takes part in `MyModule`'s `forward` calculation.
 
-See a PyTorch [example](/examples/quantization_aware_training/torch/resnet18/README.md) for **Quantization** Compression scenario on Tiny ImageNet-200 dataset.
+</details>
+
+## Examples
+
+- See a PyTorch [example](/examples/quantization_aware_training/torch/resnet18/README.md) for **Quantization** Compression scenario on Tiny ImageNet-200 dataset.
+- See a TensorFlow [example](/examples/quantization_aware_training/tensorflow/mobilenet_v2/README.md) for **Quantization** Compression scenario on `imagenette/320px-v2` dataset.

From df6bb595ad0f6b4bc71c109365c497ab0163b2c4 Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Tue, 28 Jan 2025 12:42:58 +0000
Subject: [PATCH 02/19] deprecate tf create_compressed_model

---
 nncf/tensorflow/helpers/model_creation.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/nncf/tensorflow/helpers/model_creation.py b/nncf/tensorflow/helpers/model_creation.py
index 3edbb41880e..40c04d7532e 100644
--- a/nncf/tensorflow/helpers/model_creation.py
+++ b/nncf/tensorflow/helpers/model_creation.py
@@ -18,6 +18,7 @@
 from nncf import NNCFConfig
 from nncf.api.compression import CompressionAlgorithmController
 from nncf.common.compression import BaseCompressionAlgorithmController as BaseController
+from nncf.common.deprecation import deprecated
 from nncf.common.utils.api_marker import api
 from nncf.config.extractors import extract_algorithm_names
 from nncf.config.telemetry_extractors import CompressionStartedFromConfig
@@ -62,6 +63,12 @@ def create_compression_algorithm_builder(config: NNCFConfig, should_init: bool)
         CompressionStartedFromConfig(argname="config"),
     ],
 )
+@deprecated(
+    msg="The `create_compressed_model()` method is deprecated and will be removed in a 2.14.3 release. "
+    "Consider using the 'nncf.quantize()' method instead. "
+    "Please refer to the documentation for guidance on migration.",
+    start_version="2.14.2",
+)
 def create_compressed_model(
     model: tf.keras.Model, config: NNCFConfig, compression_state: Optional[Dict[str, Any]] = None
 ) -> Tuple[CompressionAlgorithmController, tf.keras.Model]:

From 656605d5f9d054ed9ddcd2299bf6e5ce57896a71 Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Tue, 28 Jan 2025 12:57:08 +0000
Subject: [PATCH 03/19] minor update

---
 nncf/tensorflow/helpers/model_creation.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/nncf/tensorflow/helpers/model_creation.py b/nncf/tensorflow/helpers/model_creation.py
index 40c04d7532e..52b50348a49 100644
--- a/nncf/tensorflow/helpers/model_creation.py
+++ b/nncf/tensorflow/helpers/model_creation.py
@@ -64,10 +64,10 @@ def create_compression_algorithm_builder(config: NNCFConfig, should_init: bool)
     ],
 )
 @deprecated(
-    msg="The `create_compressed_model()` method is deprecated and will be removed in a 2.14.3 release. "
-    "Consider using the 'nncf.quantize()' method instead. "
+    msg="Consider using the 'nncf.quantize()' method instead. "
     "Please refer to the documentation for guidance on migration.",
     start_version="2.14.2",
+    end_version="2.14.3",
 )
 def create_compressed_model(
     model: tf.keras.Model, config: NNCFConfig, compression_state: Optional[Dict[str, Any]] = None

From 0dea99fdc859b7fff5da27138feabe7f6dc81486 Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Tue, 28 Jan 2025 13:53:29 +0000
Subject: [PATCH 04/19] update

---
 .../quantization_aware_training/Usage.md      |  2 +-
 nncf/tensorflow/helpers/model_creation.py     | 21 ++++++++++++-------
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/docs/usage/training_time_compression/quantization_aware_training/Usage.md b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
index b9f149bee82..7affc856048 100644
--- a/docs/usage/training_time_compression/quantization_aware_training/Usage.md
+++ b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
@@ -132,4 +132,4 @@ In the example above, the NNCF-compressed models that contain instances of `MyMo
 ## Examples
 
 - See a PyTorch [example](/examples/quantization_aware_training/torch/resnet18/README.md) for **Quantization** Compression scenario on Tiny ImageNet-200 dataset.
-- See a TensorFlow [example](/examples/quantization_aware_training/tensorflow/mobilenet_v2/README.md) for **Quantization** Compression scenario on `imagenette/320px-v2` dataset.
+- See a TensorFlow [example](/examples/quantization_aware_training/tensorflow/mobilenet_v2/README.md) for **Quantization** Compression scenario on imagenette/320px-v2 dataset.
diff --git a/nncf/tensorflow/helpers/model_creation.py b/nncf/tensorflow/helpers/model_creation.py
index 52b50348a49..c324811537d 100644
--- a/nncf/tensorflow/helpers/model_creation.py
+++ b/nncf/tensorflow/helpers/model_creation.py
@@ -18,7 +18,7 @@
 from nncf import NNCFConfig
 from nncf.api.compression import CompressionAlgorithmController
 from nncf.common.compression import BaseCompressionAlgorithmController as BaseController
-from nncf.common.deprecation import deprecated
+from nncf.common.deprecation import warning_deprecated
 from nncf.common.utils.api_marker import api
 from nncf.config.extractors import extract_algorithm_names
 from nncf.config.telemetry_extractors import CompressionStartedFromConfig
@@ -63,12 +63,6 @@ def create_compression_algorithm_builder(config: NNCFConfig, should_init: bool)
         CompressionStartedFromConfig(argname="config"),
     ],
 )
-@deprecated(
-    msg="Consider using the 'nncf.quantize()' method instead. "
-    "Please refer to the documentation for guidance on migration.",
-    start_version="2.14.2",
-    end_version="2.14.3",
-)
 def create_compressed_model(
     model: tf.keras.Model, config: NNCFConfig, compression_state: Optional[Dict[str, Any]] = None
 ) -> Tuple[CompressionAlgorithmController, tf.keras.Model]:
@@ -87,6 +81,19 @@ def create_compressed_model(
     :return: A tuple of the compression controller for the requested algorithm(s) and the model object with additional
      modifications necessary to enable algorithm-specific compression during fine-tuning.
     """
+
+    warning_deprecated(
+        "The 'nncf.tensorflow.create_compressed_model' function is deprecated and will be removed in a "
+        "future release.\n"
+        "To perform post training quantization (PTQ) or quantization aware training (QAT),"
+        " use the new nncf.quantize() API:\n"
+        " - https://github.com/openvinotoolkit/nncf?tab=readme-ov-file#post-training-quantization\n"
+        " - https://github.com/openvinotoolkit/nncf?tab=readme-ov-file#training-time-quantization\n"
+        "Examples:\n"
+        " - https://github.com/openvinotoolkit/nncf/tree/develop/examples/post_training_quantization/tensorflow\n"
+        " - https://github.com/openvinotoolkit/nncf/tree/develop/examples/quantization_aware_training/tensorflow"
+    )
+
     if is_experimental_quantization(config):
         if is_keras_layer_model(model):
             raise ValueError(

From adf39653c6907125339b1606857eec0ca502cdd6 Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Tue, 28 Jan 2025 14:07:32 +0000
Subject: [PATCH 05/19] update

---
 .../quantization_aware_training/Usage.md      | 30 +++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/docs/usage/training_time_compression/quantization_aware_training/Usage.md b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
index 7affc856048..2adf0806127 100644
--- a/docs/usage/training_time_compression/quantization_aware_training/Usage.md
+++ b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
@@ -104,6 +104,36 @@ You can save the `compressed_model` object `torch.save` as usual: via `state_dic
 
 </details>
 
+<details><summary><b>TensorFlow</b></summary>
+
+To save a model checkpoint, use the following API:
+
+```python
+from nncf.tensorflow.callbacks.checkpoint_callback import CheckpointManagerCallback
+
+checkpoint = tf.train.Checkpoint(model=model,
+                                 ... # the rest of the user-defined objects to save
+                                 )
+callbacks = []
+callbacks.append(CheckpointManagerCallback(checkpoint, path_to_checkpoint))
+...
+model.fit(..., callbacks=callbacks)
+```
+
+To restore the model from checkpoint, use the following API:
+
+```python
+checkpoint = tf.train.Checkpoint()
+checkpoint.restore(path_to_checkpoint)
+
+compression_ctrl, model = nncf.quantize(model, calibration_dataset)
+checkpoint = tf.train.Checkpoint(model=model,
+                                 ...)
+checkpoint.restore(path_to_checkpoint)
+```
+
+</details>
+
 ## Advanced usage
 
 ### Compression of custom modules

From 12a02cb21b3905df16927767a6dfb8a90171350b Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Thu, 30 Jan 2025 14:23:32 +0000
Subject: [PATCH 06/19] update

---
 nncf/tensorflow/__init__.py                   |  3 ++
 nncf/tensorflow/helpers/model_creation.py     | 21 +++++++++
 nncf/tensorflow/quantization/algorithm.py     | 23 ++++++----
 .../tensorflow/quantization/quantize_model.py |  5 ++-
 nncf/tensorflow/utils/state.py                | 44 ++++++++++++++++++-
 5 files changed, 85 insertions(+), 11 deletions(-)

diff --git a/nncf/tensorflow/__init__.py b/nncf/tensorflow/__init__.py
index 1bd6cef967a..091b48ac6d7 100644
--- a/nncf/tensorflow/__init__.py
+++ b/nncf/tensorflow/__init__.py
@@ -44,6 +44,7 @@
 )
 from nncf.tensorflow.helpers import create_compressed_model as create_compressed_model
 from nncf.tensorflow.helpers.callback_creation import create_compression_callbacks as create_compression_callbacks
+from nncf.tensorflow.helpers.model_creation import load_from_config
 from nncf.tensorflow.initialization import register_default_init_args as register_default_init_args
 from nncf.tensorflow.pruning.filter_pruning import algorithm as filter_pruning_algorithm
 
@@ -51,3 +52,5 @@
 from nncf.tensorflow.quantization import algorithm as quantization_algorithm
 from nncf.tensorflow.sparsity.magnitude import algorithm as magnitude_sparsity_algorithm
 from nncf.tensorflow.sparsity.rb import algorithm as rb_sparsity_algorithm
+from nncf.tensorflow.utils.state import ModelConfig
+from nncf.tensorflow.utils.state import get_config
diff --git a/nncf/tensorflow/helpers/model_creation.py b/nncf/tensorflow/helpers/model_creation.py
index c324811537d..01d46320248 100644
--- a/nncf/tensorflow/helpers/model_creation.py
+++ b/nncf/tensorflow/helpers/model_creation.py
@@ -30,8 +30,12 @@
 from nncf.tensorflow.algorithm_selector import get_compression_algorithm_builder
 from nncf.tensorflow.api.composite_compression import TFCompositeCompressionAlgorithmBuilder
 from nncf.tensorflow.api.compression import TFCompressionAlgorithmBuilder
+from nncf.tensorflow.graph.model_transformer import TFModelTransformer
+from nncf.tensorflow.graph.transformations.layout import TFTransformationLayout
 from nncf.tensorflow.graph.utils import is_keras_layer_model
 from nncf.tensorflow.helpers.utils import get_built_model
+from nncf.tensorflow.quantization.algorithm import QuantizationBuilder
+from nncf.tensorflow.utils.state import ModelConfig
 
 
 def create_compression_algorithm_builder(config: NNCFConfig, should_init: bool) -> TFCompressionAlgorithmBuilder:
@@ -140,3 +144,20 @@ def get_input_signature(config: NNCFConfig):
         input_signature.append(tf.TensorSpec(shape=shape, dtype=tf.float32))
 
     return input_signature if len(input_signature) > 1 else input_signature[0]
+
+
+def load_from_config(model: tf.keras.Model, config: ModelConfig) -> tf.keras.Model:
+    """
+    TODO(TF)
+
+    :param model:
+    :parem config:
+    :return:
+    """
+    transformation_layout = TFTransformationLayout()
+    # pylint: disable=protected-access
+    insertion_commands, _ = QuantizationBuilder._build_insertion_commands_for_quantizer_setup(config.quantizer_setup)
+    for command in insertion_commands:
+        transformation_layout.register(command)
+    model_transformer = TFModelTransformer(model)
+    return model_transformer.transform(transformation_layout)
diff --git a/nncf/tensorflow/quantization/algorithm.py b/nncf/tensorflow/quantization/algorithm.py
index 7a8a23a0ed7..3e2215991d1 100644
--- a/nncf/tensorflow/quantization/algorithm.py
+++ b/nncf/tensorflow/quantization/algorithm.py
@@ -346,14 +346,17 @@ def _get_half_range(
                 return True
         return False
 
-    def _create_quantizer(self, name: str, qspec: TFQuantizerSpec) -> Quantizer:
+    @staticmethod
+    def _create_quantizer(name: str, qspec: TFQuantizerSpec) -> Quantizer:
         quantizer_cls = NNCF_QUANTIZATION_OPERATIONS.get(qspec.mode)
         return quantizer_cls(name, qspec)
 
+    @staticmethod
     def _build_insertion_commands_for_quantizer_setup(
-        self, quantizer_setup: TFQuantizationSetup
-    ) -> List[TFInsertionCommand]:
+        quantizer_setup: TFQuantizationSetup,
+    ) -> Tuple[List[TFInsertionCommand], List[str]]:
         insertion_commands = []
+        op_names = []
         quantization_points = quantizer_setup.get_quantization_points()
         non_unified_scales_quantization_point_ids = set(range(len(quantization_points)))
 
@@ -365,7 +368,7 @@ def _build_insertion_commands_for_quantizer_setup(
             quantizer_spec = qp.quantizer_spec
             op_name = qp.op_name + "/unified_scale_group"
             quantizer = FakeQuantize(quantizer_spec, name=op_name)
-            self._op_names.append(quantizer.op_name)
+            op_names.append(quantizer.op_name)
             target_points = []
             for us_qp_id in unified_scales_group:
                 non_unified_scales_quantization_point_ids.discard(us_qp_id)
@@ -387,24 +390,26 @@ def _build_insertion_commands_for_quantizer_setup(
             quantizer_spec = quantization_point.quantizer_spec
             target_point = quantization_point.target_point
             if quantization_point.is_weight_quantization():
-                quantizer = self._create_quantizer(op_name, quantizer_spec)
-                self._op_names.append(op_name)
+                quantizer = QuantizationBuilder._create_quantizer(op_name, quantizer_spec)
+                op_names.append(op_name)
             else:
                 quantizer = FakeQuantize(quantizer_spec, name=op_name)
-                self._op_names.append(quantizer.op_name)
+                op_names.append(quantizer.op_name)
             command = TFInsertionCommand(
                 target_point=target_point,
                 callable_object=quantizer,
                 priority=TransformationPriority.QUANTIZATION_PRIORITY,
             )
             insertion_commands.append(command)
-        return insertion_commands
+        return insertion_commands, op_names
 
     def get_transformation_layout(self, model: tf.keras.Model) -> TFTransformationLayout:
         transformations = TFTransformationLayout()
         if self._quantizer_setup is None:
             self._quantizer_setup = self._get_quantizer_setup(model)
-        insertion_commands = self._build_insertion_commands_for_quantizer_setup(self._quantizer_setup)
+        insertion_commands, self._op_names = QuantizationBuilder._build_insertion_commands_for_quantizer_setup(
+            self._quantizer_setup
+        )
         for command in insertion_commands:
             transformations.register(command)
         return transformations
diff --git a/nncf/tensorflow/quantization/quantize_model.py b/nncf/tensorflow/quantization/quantize_model.py
index 02da5a28ab4..7da837850a3 100644
--- a/nncf/tensorflow/quantization/quantize_model.py
+++ b/nncf/tensorflow/quantization/quantize_model.py
@@ -176,6 +176,9 @@ def quantize_impl(
         ]
     )
 
-    _, compressed_model = create_compressed_model(model=model, config=nncf_config)
+    compression_ctrl, compressed_model = create_compressed_model(model=model, config=nncf_config)
+
+    nncf_model_config = compression_ctrl.get_compression_state()["builder_state"]["quantization"]
+    setattr(compressed_model, "_nncf_model_config", nncf_model_config)
 
     return compressed_model
diff --git a/nncf/tensorflow/utils/state.py b/nncf/tensorflow/utils/state.py
index df223c62623..8a05af52c03 100644
--- a/nncf/tensorflow/utils/state.py
+++ b/nncf/tensorflow/utils/state.py
@@ -10,11 +10,12 @@
 # limitations under the License.
 
 import json
-from typing import Any, Dict
+from typing import Any, Dict, Optional
 
 import tensorflow as tf
 
 from nncf.common.compression import BaseCompressionAlgorithmController
+from nncf.tensorflow.quantization.algorithm import TFQuantizationSetup
 
 # TODO(achurkin): remove pylint ignore after 120296 ticked is fixed
 
@@ -86,3 +87,44 @@ def deserialize(self, string_value: str) -> None:
         :param string_value: A serialized compression state.
         """
         self._state = json.loads(string_value)
+
+
+class ModelConfig(tf.train.experimental.PythonState):
+    """
+    TODO(TF)
+    """
+
+    def __init__(self, quantizer_setup: Optional[TFQuantizationSetup] = None):
+        """ """
+        self.quantizer_setup = quantizer_setup
+
+    def serialize(self) -> str:
+        """
+        Callback to serialize the model config.
+
+        :return: A serialized model config.
+        """
+        data = {"quantizer_setup": self.quantizer_setup.get_state()}
+        return json.dumps(data)
+
+    def deserialize(self, string_value: str) -> None:
+        """
+        Callback to deserialize the model config.
+
+        :param string_value: A serialized model config.
+        """
+        data = json.loads(string_value)
+        self.quantizer_setup = TFQuantizationSetup.from_state(data["quantizer_setup"])
+
+
+def get_config(model: tf.keras.Model) -> ModelConfig:
+    """
+    TODO(TF)
+
+    :param model:
+    :return:
+    """
+    data = getattr(model, "_nncf_model_config")
+    delattr(model, "_nncf_model_config")
+    quantizer_setup = TFQuantizationSetup.from_state(data["quantizer_setup"])
+    return ModelConfig(quantizer_setup)

From 9d8d4d413375ed5afb44fca4fa839b051e447cf0 Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Thu, 30 Jan 2025 16:09:25 +0000
Subject: [PATCH 07/19] update

---
 .../quantization_aware_training/Usage.md       | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/docs/usage/training_time_compression/quantization_aware_training/Usage.md b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
index 2adf0806127..5e14a3ab4f4 100644
--- a/docs/usage/training_time_compression/quantization_aware_training/Usage.md
+++ b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
@@ -109,9 +109,12 @@ You can save the `compressed_model` object `torch.save` as usual: via `state_dic
 To save a model checkpoint, use the following API:
 
 ```python
+from nncf.tensorflow import get_config
 from nncf.tensorflow.callbacks.checkpoint_callback import CheckpointManagerCallback
 
-checkpoint = tf.train.Checkpoint(model=model,
+config = get_config(quantized_model)
+checkpoint = tf.train.Checkpoint(model=quantized_model,
+                                 config=config,
                                  ... # the rest of the user-defined objects to save
                                  )
 callbacks = []
@@ -123,12 +126,17 @@ model.fit(..., callbacks=callbacks)
 To restore the model from checkpoint, use the following API:
 
 ```python
-checkpoint = tf.train.Checkpoint()
+from nncf.tensorflow import ModelConfig
+from nncf.tensorflow import load_from_config
+
+checkpoint = tf.train.Checkpoint(config=ModelConfig())
 checkpoint.restore(path_to_checkpoint)
 
-compression_ctrl, model = nncf.quantize(model, calibration_dataset)
-checkpoint = tf.train.Checkpoint(model=model,
-                                 ...)
+quantized_model = load_from_config(model, checkpoint.config)
+
+checkpoint = tf.train.Checkpoint(model=quantized_model
+                                 ... # the rest of the user-defined objects to load
+                                 )
 checkpoint.restore(path_to_checkpoint)
 ```
 

From 151ff1e0dbdc45a6978af2005467e4568f47e5d2 Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Thu, 30 Jan 2025 16:15:49 +0000
Subject: [PATCH 08/19] update

---
 .../quantization_aware_training/Usage.md                        | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/usage/training_time_compression/quantization_aware_training/Usage.md b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
index 5e14a3ab4f4..27e0de95eac 100644
--- a/docs/usage/training_time_compression/quantization_aware_training/Usage.md
+++ b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
@@ -120,7 +120,7 @@ checkpoint = tf.train.Checkpoint(model=quantized_model,
 callbacks = []
 callbacks.append(CheckpointManagerCallback(checkpoint, path_to_checkpoint))
 ...
-model.fit(..., callbacks=callbacks)
+quantized_model.fit(..., callbacks=callbacks)
 ```
 
 To restore the model from checkpoint, use the following API:

From 1ff958246664c3b61e9ce487f8edb5d6c1db8bd9 Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Thu, 30 Jan 2025 16:18:03 +0000
Subject: [PATCH 09/19] update

---
 nncf/tensorflow/helpers/model_creation.py | 2 +-
 nncf/torch/model_creation.py              | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/nncf/tensorflow/helpers/model_creation.py b/nncf/tensorflow/helpers/model_creation.py
index 01d46320248..4cc4e478625 100644
--- a/nncf/tensorflow/helpers/model_creation.py
+++ b/nncf/tensorflow/helpers/model_creation.py
@@ -90,7 +90,7 @@ def create_compressed_model(
         "The 'nncf.tensorflow.create_compressed_model' function is deprecated and will be removed in a "
         "future release.\n"
         "To perform post training quantization (PTQ) or quantization aware training (QAT),"
-        " use the new nncf.quantize() API:\n"
+        " use the nncf.quantize() API:\n"
         " - https://github.com/openvinotoolkit/nncf?tab=readme-ov-file#post-training-quantization\n"
         " - https://github.com/openvinotoolkit/nncf?tab=readme-ov-file#training-time-quantization\n"
         "Examples:\n"
diff --git a/nncf/torch/model_creation.py b/nncf/torch/model_creation.py
index 8674cb3d0ff..4d55f134fb7 100644
--- a/nncf/torch/model_creation.py
+++ b/nncf/torch/model_creation.py
@@ -106,7 +106,7 @@ def create_compressed_model(
     warning_deprecated(
         "The 'nncf.torch.create_compressed_model' function is deprecated and will be removed in a future release.\n"
         "To perform post training quantization (PTQ) or quantization aware training (QAT),"
-        " use the new nncf.quantize() API:\n"
+        " use the nncf.quantize() API:\n"
         " - https://github.com/openvinotoolkit/nncf?tab=readme-ov-file#post-training-quantization\n"
         " - https://github.com/openvinotoolkit/nncf?tab=readme-ov-file#training-time-quantization\n"
         "Examples:\n"

From e6f466c71882b6981b5788488e604d0ed01a293d Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Fri, 31 Jan 2025 14:29:57 +0000
Subject: [PATCH 10/19] update

---
 .../quantization_aware_training/Usage.md      | 11 ++++----
 nncf/tensorflow/__init__.py                   |  2 +-
 nncf/tensorflow/helpers/model_creation.py     |  9 ++++---
 nncf/tensorflow/quantization/algorithm.py     |  4 +--
 .../tensorflow/quantization/quantize_model.py |  4 +--
 nncf/tensorflow/utils/state.py                | 26 ++++++++-----------
 6 files changed, 28 insertions(+), 28 deletions(-)

diff --git a/docs/usage/training_time_compression/quantization_aware_training/Usage.md b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
index 27e0de95eac..8497dc954a4 100644
--- a/docs/usage/training_time_compression/quantization_aware_training/Usage.md
+++ b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
@@ -109,12 +109,13 @@ You can save the `compressed_model` object `torch.save` as usual: via `state_dic
 To save a model checkpoint, use the following API:
 
 ```python
+from nncf.tensorflow import ConfigState
 from nncf.tensorflow import get_config
 from nncf.tensorflow.callbacks.checkpoint_callback import CheckpointManagerCallback
 
-config = get_config(quantized_model)
+nncf_config = get_config(quantized_model)
 checkpoint = tf.train.Checkpoint(model=quantized_model,
-                                 config=config,
+                                 nncf_config_state=ConfigState(nncf_config),
                                  ... # the rest of the user-defined objects to save
                                  )
 callbacks = []
@@ -126,13 +127,13 @@ quantized_model.fit(..., callbacks=callbacks)
 To restore the model from checkpoint, use the following API:
 
 ```python
-from nncf.tensorflow import ModelConfig
+from nncf.tensorflow import ConfigState
 from nncf.tensorflow import load_from_config
 
-checkpoint = tf.train.Checkpoint(config=ModelConfig())
+checkpoint = tf.train.Checkpoint(nncf_config_state=ConfigState())
 checkpoint.restore(path_to_checkpoint)
 
-quantized_model = load_from_config(model, checkpoint.config)
+quantized_model = load_from_config(model, checkpoint.nncf_config_state.config)
 
 checkpoint = tf.train.Checkpoint(model=quantized_model
                                  ... # the rest of the user-defined objects to load
diff --git a/nncf/tensorflow/__init__.py b/nncf/tensorflow/__init__.py
index 091b48ac6d7..d060d94528b 100644
--- a/nncf/tensorflow/__init__.py
+++ b/nncf/tensorflow/__init__.py
@@ -52,5 +52,5 @@
 from nncf.tensorflow.quantization import algorithm as quantization_algorithm
 from nncf.tensorflow.sparsity.magnitude import algorithm as magnitude_sparsity_algorithm
 from nncf.tensorflow.sparsity.rb import algorithm as rb_sparsity_algorithm
-from nncf.tensorflow.utils.state import ModelConfig
+from nncf.tensorflow.utils.state import ConfigState
 from nncf.tensorflow.utils.state import get_config
diff --git a/nncf/tensorflow/helpers/model_creation.py b/nncf/tensorflow/helpers/model_creation.py
index 4cc4e478625..26da6fa822b 100644
--- a/nncf/tensorflow/helpers/model_creation.py
+++ b/nncf/tensorflow/helpers/model_creation.py
@@ -35,7 +35,7 @@
 from nncf.tensorflow.graph.utils import is_keras_layer_model
 from nncf.tensorflow.helpers.utils import get_built_model
 from nncf.tensorflow.quantization.algorithm import QuantizationBuilder
-from nncf.tensorflow.utils.state import ModelConfig
+from nncf.tensorflow.quantization.algorithm import TFQuantizationSetup
 
 
 def create_compression_algorithm_builder(config: NNCFConfig, should_init: bool) -> TFCompressionAlgorithmBuilder:
@@ -146,7 +146,7 @@ def get_input_signature(config: NNCFConfig):
     return input_signature if len(input_signature) > 1 else input_signature[0]
 
 
-def load_from_config(model: tf.keras.Model, config: ModelConfig) -> tf.keras.Model:
+def load_from_config(model: tf.keras.Model, config: Dict[str, Any]) -> tf.keras.Model:
     """
     TODO(TF)
 
@@ -154,9 +154,12 @@ def load_from_config(model: tf.keras.Model, config: ModelConfig) -> tf.keras.Mod
     :parem config:
     :return:
     """
+    quantizer_setup_state = config["quantization"]["quantizer_setup"]
+    quantizer_setup = TFQuantizationSetup.from_state(quantizer_setup_state)
+
     transformation_layout = TFTransformationLayout()
     # pylint: disable=protected-access
-    insertion_commands, _ = QuantizationBuilder._build_insertion_commands_for_quantizer_setup(config.quantizer_setup)
+    insertion_commands, _ = QuantizationBuilder.build_insertion_commands_for_quantizer_setup(quantizer_setup)
     for command in insertion_commands:
         transformation_layout.register(command)
     model_transformer = TFModelTransformer(model)
diff --git a/nncf/tensorflow/quantization/algorithm.py b/nncf/tensorflow/quantization/algorithm.py
index 3e2215991d1..d2cb5fc450d 100644
--- a/nncf/tensorflow/quantization/algorithm.py
+++ b/nncf/tensorflow/quantization/algorithm.py
@@ -352,7 +352,7 @@ def _create_quantizer(name: str, qspec: TFQuantizerSpec) -> Quantizer:
         return quantizer_cls(name, qspec)
 
     @staticmethod
-    def _build_insertion_commands_for_quantizer_setup(
+    def build_insertion_commands_for_quantizer_setup(
         quantizer_setup: TFQuantizationSetup,
     ) -> Tuple[List[TFInsertionCommand], List[str]]:
         insertion_commands = []
@@ -407,7 +407,7 @@ def get_transformation_layout(self, model: tf.keras.Model) -> TFTransformationLa
         transformations = TFTransformationLayout()
         if self._quantizer_setup is None:
             self._quantizer_setup = self._get_quantizer_setup(model)
-        insertion_commands, self._op_names = QuantizationBuilder._build_insertion_commands_for_quantizer_setup(
+        insertion_commands, self._op_names = QuantizationBuilder.build_insertion_commands_for_quantizer_setup(
             self._quantizer_setup
         )
         for command in insertion_commands:
diff --git a/nncf/tensorflow/quantization/quantize_model.py b/nncf/tensorflow/quantization/quantize_model.py
index 7da837850a3..264cd58ae4a 100644
--- a/nncf/tensorflow/quantization/quantize_model.py
+++ b/nncf/tensorflow/quantization/quantize_model.py
@@ -178,7 +178,7 @@ def quantize_impl(
 
     compression_ctrl, compressed_model = create_compressed_model(model=model, config=nncf_config)
 
-    nncf_model_config = compression_ctrl.get_compression_state()["builder_state"]["quantization"]
-    setattr(compressed_model, "_nncf_model_config", nncf_model_config)
+    config = compression_ctrl.get_compression_state()["builder_state"]["quantization"]
+    setattr(compressed_model, "_nncf_config", config)
 
     return compressed_model
diff --git a/nncf/tensorflow/utils/state.py b/nncf/tensorflow/utils/state.py
index 8a05af52c03..5306df5aaf3 100644
--- a/nncf/tensorflow/utils/state.py
+++ b/nncf/tensorflow/utils/state.py
@@ -15,7 +15,6 @@
 import tensorflow as tf
 
 from nncf.common.compression import BaseCompressionAlgorithmController
-from nncf.tensorflow.quantization.algorithm import TFQuantizationSetup
 
 # TODO(achurkin): remove pylint ignore after 120296 ticked is fixed
 
@@ -89,23 +88,22 @@ def deserialize(self, string_value: str) -> None:
         self._state = json.loads(string_value)
 
 
-class ModelConfig(tf.train.experimental.PythonState):
+class ConfigState(tf.train.experimental.PythonState):
     """
     TODO(TF)
     """
 
-    def __init__(self, quantizer_setup: Optional[TFQuantizationSetup] = None):
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
         """ """
-        self.quantizer_setup = quantizer_setup
+        self.config = config
 
     def serialize(self) -> str:
         """
-        Callback to serialize the model config.
+        Callback to serialize the config.
 
-        :return: A serialized model config.
+        :return: A serialized config.
         """
-        data = {"quantizer_setup": self.quantizer_setup.get_state()}
-        return json.dumps(data)
+        return json.dumps(self.config)
 
     def deserialize(self, string_value: str) -> None:
         """
@@ -113,18 +111,16 @@ def deserialize(self, string_value: str) -> None:
 
         :param string_value: A serialized model config.
         """
-        data = json.loads(string_value)
-        self.quantizer_setup = TFQuantizationSetup.from_state(data["quantizer_setup"])
+        self.config = json.loads(string_value)
 
 
-def get_config(model: tf.keras.Model) -> ModelConfig:
+def get_config(model: tf.keras.Model) -> Dict[str, Any]:
     """
     TODO(TF)
 
     :param model:
     :return:
     """
-    data = getattr(model, "_nncf_model_config")
-    delattr(model, "_nncf_model_config")
-    quantizer_setup = TFQuantizationSetup.from_state(data["quantizer_setup"])
-    return ModelConfig(quantizer_setup)
+    config = getattr(model, "_nncf_config")
+    delattr(model, "_nncf_config")
+    return config

From eda953e8ceb82c1bd014cf68185679e91a956c93 Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Fri, 31 Jan 2025 15:00:02 +0000
Subject: [PATCH 11/19] update

---
 .../tensorflow/mobilenet_v2/main.py              |  4 ++--
 nncf/tensorflow/helpers/model_creation.py        | 16 ++++++++++++----
 nncf/tensorflow/quantization/quantize_model.py   |  3 +++
 nncf/tensorflow/utils/state.py                   | 12 +++++++-----
 4 files changed, 24 insertions(+), 11 deletions(-)

diff --git a/examples/quantization_aware_training/tensorflow/mobilenet_v2/main.py b/examples/quantization_aware_training/tensorflow/mobilenet_v2/main.py
index cf3bc372887..233ec512727 100644
--- a/examples/quantization_aware_training/tensorflow/mobilenet_v2/main.py
+++ b/examples/quantization_aware_training/tensorflow/mobilenet_v2/main.py
@@ -167,8 +167,8 @@ def transform_fn(data_item):
 ###############################################################################
 # Benchmark performance, calculate compression rate and validate accuracy
 
-ov_model = ov.convert_model(tf_model, share_weights=False)
-ov_quantized_model = ov.convert_model(stripped_model, share_weights=False)
+ov_model = ov.convert_model(tf_model)
+ov_quantized_model = ov.convert_model(stripped_model)
 
 fp32_ir_path = ROOT / "mobilenet_v2_fp32.xml"
 ov.save_model(ov_model, fp32_ir_path, compress_to_fp16=False)
diff --git a/nncf/tensorflow/helpers/model_creation.py b/nncf/tensorflow/helpers/model_creation.py
index 26da6fa822b..ac3201aeeb5 100644
--- a/nncf/tensorflow/helpers/model_creation.py
+++ b/nncf/tensorflow/helpers/model_creation.py
@@ -25,6 +25,7 @@
 from nncf.config.utils import is_experimental_quantization
 from nncf.telemetry import tracked_function
 from nncf.telemetry.events import NNCF_TF_CATEGORY
+from nncf.telemetry.extractors import FunctionCallTelemetryExtractor
 from nncf.tensorflow.accuracy_aware_training.keras_model_utils import accuracy_aware_fit
 from nncf.tensorflow.algorithm_selector import NoCompressionAlgorithmBuilder
 from nncf.tensorflow.algorithm_selector import get_compression_algorithm_builder
@@ -146,13 +147,20 @@ def get_input_signature(config: NNCFConfig):
     return input_signature if len(input_signature) > 1 else input_signature[0]
 
 
+@tracked_function(
+    NNCF_TF_CATEGORY,
+    [
+        FunctionCallTelemetryExtractor("nncf.tensorflow.load_from_config"),
+    ],
+)
 def load_from_config(model: tf.keras.Model, config: Dict[str, Any]) -> tf.keras.Model:
     """
-    TODO(TF)
+    Recovers additional modules from given config.
+    Does not recover additional modules weights as they are located in a corresponded checkpoint file.
 
-    :param model:
-    :parem config:
-    :return:
+    :param model: TensorFlow model.
+    :parem config: Config.
+    :return: tf.keras.Model builded from given model with additional layers recovered from given config.
     """
     quantizer_setup_state = config["quantization"]["quantizer_setup"]
     quantizer_setup = TFQuantizationSetup.from_state(quantizer_setup_state)
diff --git a/nncf/tensorflow/quantization/quantize_model.py b/nncf/tensorflow/quantization/quantize_model.py
index 264cd58ae4a..9295d838aea 100644
--- a/nncf/tensorflow/quantization/quantize_model.py
+++ b/nncf/tensorflow/quantization/quantize_model.py
@@ -178,6 +178,9 @@ def quantize_impl(
 
     compression_ctrl, compressed_model = create_compressed_model(model=model, config=nncf_config)
 
+    # NOTE: We set the config here to properly save/load the quantized model during training into tf.train.Checkpoint.
+    # You can obtain that config via the nncf.tensorflow.get_config() method and save/load it to/from
+    # tf.train.Checkpoint using the nncf.tensorflow.ConfigState class.
     config = compression_ctrl.get_compression_state()["builder_state"]["quantization"]
     setattr(compressed_model, "_nncf_config", config)
 
diff --git a/nncf/tensorflow/utils/state.py b/nncf/tensorflow/utils/state.py
index 5306df5aaf3..8bc3195d244 100644
--- a/nncf/tensorflow/utils/state.py
+++ b/nncf/tensorflow/utils/state.py
@@ -90,11 +90,13 @@ def deserialize(self, string_value: str) -> None:
 
 class ConfigState(tf.train.experimental.PythonState):
     """
-    TODO(TF)
+    Used to save/load a config into the tf.train.Checkpoint.
     """
 
     def __init__(self, config: Optional[Dict[str, Any]] = None):
-        """ """
+        """
+        :param config: Config.
+        """
         self.config = config
 
     def serialize(self) -> str:
@@ -116,10 +118,10 @@ def deserialize(self, string_value: str) -> None:
 
 def get_config(model: tf.keras.Model) -> Dict[str, Any]:
     """
-    TODO(TF)
+    Extracts the config from the model.
 
-    :param model:
-    :return:
+    :param model: Model.
+    :return: Config.
     """
     config = getattr(model, "_nncf_config")
     delattr(model, "_nncf_config")

From 6024dce4dfa70baed12aa2090ea44a0ed70cbaf7 Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Fri, 31 Jan 2025 15:09:26 +0000
Subject: [PATCH 12/19] update

---
 nncf/tensorflow/quantization/quantize_model.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/nncf/tensorflow/quantization/quantize_model.py b/nncf/tensorflow/quantization/quantize_model.py
index 9295d838aea..7a5e1b7798a 100644
--- a/nncf/tensorflow/quantization/quantize_model.py
+++ b/nncf/tensorflow/quantization/quantize_model.py
@@ -9,6 +9,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import warnings
 from typing import Any, Dict, Optional
 
 import tensorflow as tf
@@ -176,7 +177,9 @@ def quantize_impl(
         ]
     )
 
+    warnings.filterwarnings("ignore", category=FutureWarning)
     compression_ctrl, compressed_model = create_compressed_model(model=model, config=nncf_config)
+    warnings.filterwarnings("default", category=FutureWarning)
 
     # NOTE: We set the config here to properly save/load the quantized model during training into tf.train.Checkpoint.
     # You can obtain that config via the nncf.tensorflow.get_config() method and save/load it to/from

From 1c51fd8761d6cffd3f4d2c6eb17b65845e0744b8 Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Fri, 31 Jan 2025 19:49:07 +0000
Subject: [PATCH 13/19] update

---
 nncf/tensorflow/utils/state.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/nncf/tensorflow/utils/state.py b/nncf/tensorflow/utils/state.py
index 8bc3195d244..35578fe1c91 100644
--- a/nncf/tensorflow/utils/state.py
+++ b/nncf/tensorflow/utils/state.py
@@ -15,6 +15,7 @@
 import tensorflow as tf
 
 from nncf.common.compression import BaseCompressionAlgorithmController
+from nncf.tensorflow.quantization.algorithm import TFQuantizationSetup
 
 # TODO(achurkin): remove pylint ignore after 120296 ticked is fixed
 
@@ -105,7 +106,12 @@ def serialize(self) -> str:
 
         :return: A serialized config.
         """
-        return json.dumps(self.config)
+        data = {
+            "quantization": {
+                "quantizer_setup": TFQuantizationSetup.from_state(self.config["quantizer_setup"]).get_state(),
+            }
+        }
+        return json.dumps(data)
 
     def deserialize(self, string_value: str) -> None:
         """

From 84896d77b949ce61334a6d3988215a58ae825b4a Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Fri, 31 Jan 2025 20:13:16 +0000
Subject: [PATCH 14/19] update

---
 nncf/tensorflow/quantization/quantize_model.py | 2 +-
 nncf/tensorflow/utils/state.py                 | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/nncf/tensorflow/quantization/quantize_model.py b/nncf/tensorflow/quantization/quantize_model.py
index 7a5e1b7798a..9030ac6e276 100644
--- a/nncf/tensorflow/quantization/quantize_model.py
+++ b/nncf/tensorflow/quantization/quantize_model.py
@@ -184,7 +184,7 @@ def quantize_impl(
     # NOTE: We set the config here to properly save/load the quantized model during training into tf.train.Checkpoint.
     # You can obtain that config via the nncf.tensorflow.get_config() method and save/load it to/from
     # tf.train.Checkpoint using the nncf.tensorflow.ConfigState class.
-    config = compression_ctrl.get_compression_state()["builder_state"]["quantization"]
+    config = compression_ctrl.get_compression_state()["builder_state"]
     setattr(compressed_model, "_nncf_config", config)
 
     return compressed_model
diff --git a/nncf/tensorflow/utils/state.py b/nncf/tensorflow/utils/state.py
index 35578fe1c91..51a8c6cfa78 100644
--- a/nncf/tensorflow/utils/state.py
+++ b/nncf/tensorflow/utils/state.py
@@ -106,9 +106,10 @@ def serialize(self) -> str:
 
         :return: A serialized config.
         """
+        quantizer_setup_state = self.config["quantization"]["quantizer_setup"]
         data = {
             "quantization": {
-                "quantizer_setup": TFQuantizationSetup.from_state(self.config["quantizer_setup"]).get_state(),
+                "quantizer_setup": TFQuantizationSetup.from_state(quantizer_setup_state).get_state(),
             }
         }
         return json.dumps(data)

From f87c96aea4e37bd88a2eeef01d79301b33f8d1b0 Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Mon, 3 Feb 2025 21:30:55 +0000
Subject: [PATCH 15/19] update

---
 .../quantization_aware_training/Usage.md                        | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/usage/training_time_compression/quantization_aware_training/Usage.md b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
index 8497dc954a4..8ae7163e677 100644
--- a/docs/usage/training_time_compression/quantization_aware_training/Usage.md
+++ b/docs/usage/training_time_compression/quantization_aware_training/Usage.md
@@ -64,7 +64,7 @@ import openvino as ov
 # resulting in a clean, fully quantized model ready for deployment.
 stripped_model = nncf.strip(quantized_model)
 
-ov_quantized_model = ov.convert_model(stripped_model, share_weights=False)
+ov_quantized_model = ov.convert_model(stripped_model)
 ```
 
 </details>

From fed41ca9e9e8d4832f19b7a45b2c0c4cfc84786c Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Tue, 4 Feb 2025 11:47:04 +0000
Subject: [PATCH 16/19] update

---
 .../tensorflow/mobilenet_v2/main.py                           | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/examples/post_training_quantization/tensorflow/mobilenet_v2/main.py b/examples/post_training_quantization/tensorflow/mobilenet_v2/main.py
index 8e175f7dd3f..5f22d516e22 100644
--- a/examples/post_training_quantization/tensorflow/mobilenet_v2/main.py
+++ b/examples/post_training_quantization/tensorflow/mobilenet_v2/main.py
@@ -151,8 +151,8 @@ def transform_fn(data_item):
 ###############################################################################
 # Benchmark performance, calculate compression rate and validate accuracy
 
-ov_model = ov.convert_model(tf_model, share_weights=False)
-ov_quantized_model = ov.convert_model(tf_quantized_model, share_weights=False)
+ov_model = ov.convert_model(tf_model)
+ov_quantized_model = ov.convert_model(tf_quantized_model)
 
 fp32_ir_path = ROOT / "mobilenet_v2_fp32.xml"
 ov.save_model(ov_model, fp32_ir_path, compress_to_fp16=False)

From e317038d9bab9ac4e1b3a6c76493a67165d7c368 Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Tue, 4 Feb 2025 17:49:10 +0000
Subject: [PATCH 17/19] Add create_compressed_model_impl

---
 nncf/tensorflow/helpers/model_creation.py      | 8 ++++++++
 nncf/tensorflow/quantization/quantize_model.py | 7 ++-----
 2 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/nncf/tensorflow/helpers/model_creation.py b/nncf/tensorflow/helpers/model_creation.py
index ac3201aeeb5..b11b2ab65ee 100644
--- a/nncf/tensorflow/helpers/model_creation.py
+++ b/nncf/tensorflow/helpers/model_creation.py
@@ -98,7 +98,15 @@ def create_compressed_model(
         " - https://github.com/openvinotoolkit/nncf/tree/develop/examples/post_training_quantization/tensorflow\n"
         " - https://github.com/openvinotoolkit/nncf/tree/develop/examples/quantization_aware_training/tensorflow"
     )
+    return create_compressed_model_impl(model, config, compression_state)
 
+
+def create_compressed_model_impl(
+    model: tf.keras.Model, config: NNCFConfig, compression_state: Optional[Dict[str, Any]] = None
+) -> Tuple[CompressionAlgorithmController, tf.keras.Model]:
+    """
+    Implementation of the create_compressed_model() method.
+    """
     if is_experimental_quantization(config):
         if is_keras_layer_model(model):
             raise ValueError(
diff --git a/nncf/tensorflow/quantization/quantize_model.py b/nncf/tensorflow/quantization/quantize_model.py
index 9030ac6e276..20d164070bb 100644
--- a/nncf/tensorflow/quantization/quantize_model.py
+++ b/nncf/tensorflow/quantization/quantize_model.py
@@ -9,7 +9,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import warnings
 from typing import Any, Dict, Optional
 
 import tensorflow as tf
@@ -29,7 +28,7 @@
 from nncf.quantization.advanced_parameters import apply_advanced_parameters_to_config
 from nncf.scopes import IgnoredScope
 from nncf.scopes import convert_ignored_scope_to_list
-from nncf.tensorflow.helpers.model_creation import create_compressed_model
+from nncf.tensorflow.helpers.model_creation import create_compressed_model_impl
 
 DEFAULT_RANGE_TYPE = "mean_min_max"
 
@@ -177,9 +176,7 @@ def quantize_impl(
         ]
     )
 
-    warnings.filterwarnings("ignore", category=FutureWarning)
-    compression_ctrl, compressed_model = create_compressed_model(model=model, config=nncf_config)
-    warnings.filterwarnings("default", category=FutureWarning)
+    compression_ctrl, compressed_model = create_compressed_model_impl(model=model, config=nncf_config)
 
     # NOTE: We set the config here to properly save/load the quantized model during training into tf.train.Checkpoint.
     # You can obtain that config via the nncf.tensorflow.get_config() method and save/load it to/from

From 269cfa9779d35a3947151eda7c462b7901153450 Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Tue, 4 Feb 2025 17:54:14 +0000
Subject: [PATCH 18/19] Move get_config

---
 nncf/tensorflow/__init__.py               |  2 +-
 nncf/tensorflow/helpers/model_creation.py | 18 ++++++++++++++++++
 nncf/tensorflow/utils/state.py            | 12 ------------
 3 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/nncf/tensorflow/__init__.py b/nncf/tensorflow/__init__.py
index d060d94528b..cc71a9d6c91 100644
--- a/nncf/tensorflow/__init__.py
+++ b/nncf/tensorflow/__init__.py
@@ -44,6 +44,7 @@
 )
 from nncf.tensorflow.helpers import create_compressed_model as create_compressed_model
 from nncf.tensorflow.helpers.callback_creation import create_compression_callbacks as create_compression_callbacks
+from nncf.tensorflow.helpers.model_creation import get_config
 from nncf.tensorflow.helpers.model_creation import load_from_config
 from nncf.tensorflow.initialization import register_default_init_args as register_default_init_args
 from nncf.tensorflow.pruning.filter_pruning import algorithm as filter_pruning_algorithm
@@ -53,4 +54,3 @@
 from nncf.tensorflow.sparsity.magnitude import algorithm as magnitude_sparsity_algorithm
 from nncf.tensorflow.sparsity.rb import algorithm as rb_sparsity_algorithm
 from nncf.tensorflow.utils.state import ConfigState
-from nncf.tensorflow.utils.state import get_config
diff --git a/nncf/tensorflow/helpers/model_creation.py b/nncf/tensorflow/helpers/model_creation.py
index b11b2ab65ee..99b4aa8faac 100644
--- a/nncf/tensorflow/helpers/model_creation.py
+++ b/nncf/tensorflow/helpers/model_creation.py
@@ -180,3 +180,21 @@ def load_from_config(model: tf.keras.Model, config: Dict[str, Any]) -> tf.keras.
         transformation_layout.register(command)
     model_transformer = TFModelTransformer(model)
     return model_transformer.transform(transformation_layout)
+
+
+@tracked_function(
+    NNCF_TF_CATEGORY,
+    [
+        FunctionCallTelemetryExtractor("nncf.tensorflow.get_config"),
+    ],
+)
+def get_config(model: tf.keras.Model) -> Dict[str, Any]:
+    """
+    Extracts the config from the model.
+
+    :param model: Model.
+    :return: Config.
+    """
+    config = getattr(model, "_nncf_config")
+    delattr(model, "_nncf_config")
+    return config
diff --git a/nncf/tensorflow/utils/state.py b/nncf/tensorflow/utils/state.py
index 51a8c6cfa78..1cac500579a 100644
--- a/nncf/tensorflow/utils/state.py
+++ b/nncf/tensorflow/utils/state.py
@@ -121,15 +121,3 @@ def deserialize(self, string_value: str) -> None:
         :param string_value: A serialized model config.
         """
         self.config = json.loads(string_value)
-
-
-def get_config(model: tf.keras.Model) -> Dict[str, Any]:
-    """
-    Extracts the config from the model.
-
-    :param model: Model.
-    :return: Config.
-    """
-    config = getattr(model, "_nncf_config")
-    delattr(model, "_nncf_config")
-    return config

From 525799f83b64603d52dc4c4e482b385abcc88845 Mon Sep 17 00:00:00 2001
From: Andrey Churkin <andrey.churkin@intel.com>
Date: Tue, 4 Feb 2025 18:21:02 +0000
Subject: [PATCH 19/19] update

---
 nncf/tensorflow/helpers/model_creation.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/nncf/tensorflow/helpers/model_creation.py b/nncf/tensorflow/helpers/model_creation.py
index 99b4aa8faac..b3c506edac4 100644
--- a/nncf/tensorflow/helpers/model_creation.py
+++ b/nncf/tensorflow/helpers/model_creation.py
@@ -196,5 +196,4 @@ def get_config(model: tf.keras.Model) -> Dict[str, Any]:
     :return: Config.
     """
     config = getattr(model, "_nncf_config")
-    delattr(model, "_nncf_config")
     return config