update docs

Signed-off-by: changwangss <[email protected]>

docs/distillation.md

@@ -49,39 +49,19 @@ Where $D$ is a distance measurement as before, $F_t^{n_i}$ the output feature of
 ## usage
 ### Pytorch Script:
 ```python
-from intel_extension_for_transformers.transformers import metric, objectives, DistillationConfig, Criterion
+
 from intel_extension_for_transformers.transformers.trainer import NLPTrainer
+from neural_compressor.config import DistillationConfig
 # Replace transformers.Trainer with NLPTrainer
 # trainer = transformers.Trainer(......)
 trainer = NLPTrainer(......)
 metric = metrics.Metric(name="eval_accuracy")
-d_conf = DistillationConfig(metrics=tune_metric)
-model = trainer.distill(
-    distillation_config=d_conf, teacher_model=teacher_model
-)
+d_conf = DistillationConfig(teacher_model=teacher_model, criterion=criterion)
+model = trainer.distill(distillation_config=d_conf)
 ```
 
 Please refer to [example](../examples/huggingface/pytorch/text-classification/distillation/run_glue.py) for the details.
 
-### Tensorflow Script:
-```python
-from intel_extension_for_transformers.transformers import (DistillationConfig, metrics)
-from intel_extension_for_transformers.transformers.distillation import Criterion
-
-optimizer = TFOptimization(...)
-metric_ = metrics.Metric(name="eval_accuracy")
-criterion = Criterion(name='KnowledgeLoss',
-                    layer_mappings=[['classifier', 'classifier']],
-                    loss_types=['CE', 'CE'],
-                    loss_weight_ratio=[0.5, 0.5],
-                    add_origin_loss=False)
-distillation_conf = DistillationConfig(metrics=metric_,
-                                        criterion=criterion)
-distilled_model = optimizer.distill(
-            distillation_config=distillation_conf,
-            teacher_model=teacher_model)
-```
-Please refer to [example](../examples/huggingface/tensorflow/text-classification/distillation/run_glue.py) for the details.
 ### Create an Instance of Metric
 The Metric defines which metric will be used to measure the performance of tuned models.
 - example:
@@ -94,19 +74,23 @@ The Metric defines which metric will be used to measure the performance of tuned
 ### Create an Instance of Criterion(Optional)
 The criterion used in training phase.
 
-- arguments:
+- KnowledgeDistillationLossConfig arguments:
     |Argument   |Type       |Description                                        |Default value    |
     |:----------|:----------|:-----------------------------------------------|:----------------|
-    |name       |String|Name of criterion, like:"KnowledgeLoss", "IntermediateLayersLoss"  |"KnowledgeLoss"|
     |temperature|Float |parameter for KnowledgeDistillationLoss               |1.0             |
     |loss_types|List of string|Type of loss                               |['CE', 'CE']        |
     |loss_weight_ratio|List of float|weight ratio of loss                 |[0.5, 0.5]     |
+
+- IntermediateLayersKnowledgeDistillationLossConfig arguments:
+    |Argument   |Type       |Description                                        |Default value    |
+    |:----------|:----------|:-----------------------------------------------|:----------------|
+    |loss_types|List of string|Type of loss                               |['CE', 'CE']        |
+    |loss_weight_ratio|List of float|weight ratio of loss                 |[0.5, 0.5]     |
     |layer_mappings|List|parameter for IntermediateLayersLoss             |[] |
     |add_origin_loss|bool|parameter for IntermediateLayersLoss            |False |
-
 - example:
     ```python
-    criterion = Criterion(name='KnowledgeLoss')
+    criterion = KnowledgeDistillationLossConfig()
     ```
 
 ### Create an Instance of DistillationConfig
@@ -115,20 +99,18 @@ The DistillationConfig contains all the information related to the model distill
 - arguments:
     |Argument   |Type       |Description                                        |Default value    |
     |:----------|:----------|:-----------------------------------------------|:----------------|
-    |framework  |string     |which framework you used                        |"pytorch"        |
-    |criterion|Criterion |criterion of training                              |"KnowledgeLoss"|
-    |metrics    |Metric    |Used to evaluate accuracy of tuning model, no need for NoTrainerOptimizer|None    |
+    |teacher_model  |torch.nn.Module     | teacher model object                    |None        |
+    |criterion|Criterion |criterion of training                              |KnowledgeLoss object|
+
 
 - example:
     ```python
-    d_conf = DistillationConfig(metrics=metric, criterion=criterion)
+    d_conf = DistillationConfig(teacher_model=teacher_model, criterion=criterion)
     ```
 
 ### Distill with Trainer
 - Distill with Trainer
     NLPTrainer inherits from transformers.Trainer, so you can create a trainer as in examples of Transformers. Then you can distill model with trainer.distill function.
     ```python
-    model = trainer.distill(
-        distillation_config=d_conf, teacher_model=teacher_model
-    )
+    model = trainer.distill(distillation_config=d_conf)
     ```

docs/pruning.md

@@ -45,7 +45,7 @@ The Metric defines which metric will be used to measure the performance of tuned
 
     Please refer to [metrics document](metrics.md) for the details.
 
-### Create list of an instance of PrunerConfig(Optional)
+### Create an instance of WeightPruningConfig
 PrunerConfig defines which pruning algorithm to use and how to apply it during the training process. Intel® Extension for Transformers supports pruning types "BasicMagnitude", "PatternLock", and "GroupLasso". You can create different pruners for different layers.
 
 - arguments:
@@ -60,26 +60,18 @@ PrunerConfig defines which pruning algorithm to use and how to apply it during t
     |names|list of string|List of weight name to be pruned. If no weight is specified, all weights of the model will be pruned|[]|
     |parameters|dict of string|The hyper-parameters for pruning, refer to [the link](https://github.com/intel/neural-compressor/blob/master/docs/source/pruning.md)|None|
 
-- example:
-    ```python
-    pruner_config = PrunerConfig(prune_type='BasicMagnitude', target_sparsity_ratio=0.9)
-    ```
-
-### Create an instance of PruningConfig
-The PruningConfig contains all the information related to the model pruning behavior. If you have created Metric and PrunerConfig instance, then you can create an instance of PruningConfig. Metric and pruner are optional.
-
-- arguments:
-    |Argument   |Type       |Description                                        |Default value    |
-    |:----------|:----------|:-----------------------------------------------|:----------------|
-    |framework  |string     |Which framework you used                        |"pytorch"        |
-    |initial_sparsity_ratio|float |Initial sparsity goal, if pruner_config argument is defined, it didn't need                       |0.0|
-    |target_sparsity_ratio|float |Target sparsity goal, if pruner argument is defined, it didn't need                       |0.97|
-    |metrics    |Metric    |Used to evaluate accuracy of tuning model, no need for NoTrainerOptimizer|None    |
-    |pruner_config |PrunerConfig    |Defined pruning behavior, if it is None, then NLP will create a default a pruner with 'BasicMagnitude' pruning type                                  |None              |
+The WeightPruningConfig contains all the information related to the model pruning behavior. If you have created Metric and WeightPruningConfig instance, then you can create an instance of WeightPruningConfig. Metric and pruner are optional.
 
 - example:
     ```python
-    pruning_conf = PruningConfig(pruner_config=[pruner_config], metrics=tune_metric)
+    from neural_compressor.config import WeightPruningConfig
+
+    metric = metrics.Metric(name="eval_accuracy")
+    trainer.metrics = tune_metric
+    pruning_conf = WeightPruningConfig([{"start_step": 0, "end_step": 2}],
+                                        target_sparsity=0,9,
+                                        pruning_scope="local",
+                                        pruning_type="magnitude")
     ```
 
 ### Prune with Trainer

examples/huggingface/pytorch/question-answering/pruning/basic_magnitude/run_qa.py

@@ -26,7 +26,8 @@
 import transformers
 from dataclasses import dataclass, field
 from datasets import load_dataset, load_metric
-from intel_extension_for_transformers.transformers import metrics, OptimizedModel, PrunerConfig, PruningConfig, PruningMode
+from intel_extension_for_transformers.transformers import metrics, OptimizedModel
+from neural_compressor.config import WeightPruningConfig
 from trainer_qa import QuestionAnsweringTrainer
 from transformers import (
     AutoConfig,
@@ -210,8 +211,8 @@ class OptimizationArguments:
         metadata={"help": "Whether or not to apply prune."},
     )
     pruning_approach: Optional[str] = field(
-        default="BasicMagnitude",
-        metadata={"help": "Pruning approach. Supported approach is basic_magnite."},
+        default="magnitude",
+        metadata={"help": "Pruning approach. Supported approach is magnite."},
     )
     target_sparsity_ratio: Optional[float] = field(
         default=None,
@@ -631,12 +632,15 @@ def compute_metrics(p: EvalPrediction):
             raise ValueError("do_train must be set to True for pruning.")
 
         tune_metric = metrics.Metric(name=metric_name)
-        prune_type = 'BasicMagnitude' if optim_args.pruning_approach else optim_args.pruning_approach
+        prune_type = optim_args.pruning_approach \
+            if optim_args.pruning_approach else 'pattern_lock'
         target_sparsity_ratio = optim_args.target_sparsity_ratio \
             if optim_args.target_sparsity_ratio else None
-        pruner_config = PrunerConfig(prune_type=prune_type, target_sparsity_ratio=target_sparsity_ratio)
-        pruning_conf = PruningConfig(pruner_config=pruner_config, metrics=tune_metric)
-
+        trainer.metrics = tune_metric
+        pruning_conf = WeightPruningConfig([{"start_step": 0, "end_step": 2}],
+                                            target_sparsity=target_sparsity_ratio,
+                                            pruning_scope="local",
+                                            pruning_type=prune_type)
         model = trainer.prune(pruning_config=pruning_conf)
         trainer.save_model(training_args.output_dir)
 

examples/huggingface/pytorch/text-classification/pruning/run_glue.py

@@ -29,9 +29,8 @@
 from intel_extension_for_transformers.transformers import (
     metrics,
     OptimizedModel,
-    PrunerConfig,
-    PruningConfig,
 )
+from neural_compressor.config import WeightPruningConfig
 from intel_extension_for_transformers.transformers.trainer import NLPTrainer
 from transformers import (
     AutoConfig,
@@ -47,7 +46,6 @@
 )
 from transformers.trainer_utils import get_last_checkpoint
 from transformers.utils import check_min_version
-from transformers.utils.fx import symbolic_trace
 from typing import Optional
 
 
@@ -204,8 +202,8 @@ class OptimizationArguments:
         metadata={"help": "Whether or not to apply prune."},
     )
     pruning_approach: Optional[str] = field(
-        default="BasicMagnitude",
-        metadata={"help": "Pruning approach. Supported approach is basic_magnite."},
+        default="magnitude",
+        metadata={"help": "Pruning approach. Supported approach is magnite."},
     )
     target_sparsity_ratio: Optional[float] = field(
         default=None,
@@ -521,13 +519,15 @@ def compute_metrics(p: EvalPrediction):
             raise ValueError("do_train must be set to True for pruning.")
 
         tune_metric = metrics.Metric(name=metric_name)
-        prune_type = 'BasicMagnitude' \
-            if optim_args.pruning_approach else optim_args.pruning_approach
+        prune_type = optim_args.pruning_approach \
+            if optim_args.pruning_approach else 'pattern_lock'
         target_sparsity_ratio = optim_args.target_sparsity_ratio \
             if optim_args.target_sparsity_ratio else None
-        pruner_config = PrunerConfig(prune_type=prune_type, target_sparsity_ratio=target_sparsity_ratio)
-        pruning_conf = PruningConfig(pruner_config=pruner_config, metrics=tune_metric)
-
+        trainer.metrics = tune_metric
+        pruning_conf = WeightPruningConfig([{"start_step": 0, "end_step": 2}],
+                                            target_sparsity=target_sparsity_ratio,
+                                            pruning_scope="local",
+                                            pruning_type=prune_type)
         model = trainer.prune(pruning_config=pruning_conf)
         trainer.save_model(training_args.output_dir)
 

intel_extension_for_transformers/transformers/config.py

@@ -18,16 +18,11 @@
 
 import yaml
 from enum import Enum
-from neural_compressor.conf.config import (
-    Distillation_Conf, Pruner, Pruning_Conf, Quantization_Conf
-)
-from neural_compressor.conf.dotdict import DotDict, deep_set
+
+from neural_compressor.conf.dotdict import DotDict
 from .utils.metrics import Metric
 from .utils.objectives import Objective, performance
-from .quantization import QuantizationMode, SUPPORTED_QUANT_MODE
-from .distillation import (
-    Criterion, DistillationCriterionMode, SUPPORTED_DISTILLATION_CRITERION_MODE
-)
+
 from typing import List, Union
 from xmlrpc.client import boolean