Skip to content

Commit

Permalink
refactor: Make a tokenizer during model preparation
Browse files Browse the repository at this point in the history
Needed to tokenize example input for torch model later
  • Loading branch information
AdiKsOnDev committed Apr 24, 2024
1 parent 512aa63 commit 0041998
Showing 1 changed file with 7 additions and 19 deletions.
26 changes: 7 additions & 19 deletions tests/post_training/pipelines/lm_weight_compression.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,11 +18,12 @@

import numpy as np
import openvino as ov
import transformers
from datasets import load_dataset
from memory_profiler import memory_usage
from nncf.parameters import CompressWeightsMode
from optimum.intel.openvino import OVModelForCausalLM
from transformers import AutoTokenizer, AutoModelForCausalLM
from transformers import AutoModelForCausalLM
from transformers import AutoTokenizer
from whowhatbench import Evaluator

import nncf
Expand Down Expand Up @@ -78,16 +79,16 @@ def prepare_model(self) -> None:
self.MODEL_NAME = "torch_model.xml"
self.MODEL_FUNC = AutoModelForCausalLM
MODEL_SPECIFIC_PARAMS = {}
self.tokenizer = transformers.AutoTokenizer.from_pretrained(self.model_id)
else:
MODEL_SPECIFIC_PARAMS = {"export": True, "compile": False, "stateful": is_stateful}
if is_stateful:
self.fp32_model_dir = self.fp32_model_dir.parent / (self.fp32_model_dir.name + "_sf")
if not (self.fp32_model_dir / self.MODEL_NAME).exists():
# export by model_id
self.model_hf = self.MODEL_FUNC.from_pretrained(
self.model_id, load_in_8bit=False, **MODEL_SPECIFIC_PARAMS
)
self._dump_model_fp32()
self.model_hf = self.MODEL_FUNC.from_pretrained(self.model_id, load_in_8bit=False, **MODEL_SPECIFIC_PARAMS)
if self.backend != BackendType.TORCH:
self._dump_model_fp32()
else:
# no export, load from IR. Applicable for sequential run of test cases in local environment.
self.model_hf = self.MODEL_FUNC.from_pretrained(
Expand Down Expand Up @@ -182,19 +183,6 @@ def _dump_model_fp32(self) -> None:
self.model_hf.save_pretrained(self.fp32_model_dir)
self.model_hf._save_config(self.fp32_model_dir)

# def _compress(self):
# """
# Actual call of weight compression
# """
# if self.compression_params["mode"] == CompressWeightsMode.INT8_ASYM:
# """If compression mode is INT8, don't use a dataset as it's Unsupported"""
# self.calibration_dataset = None
#
# self.compressed_model = nncf.compress_weights(
# self.model,
# dataset=self.calibration_dataset,
# **self.compression_params,
# )
def _compress(self):
"""
Actual call of weight compression
Expand Down

0 comments on commit 0041998

Please sign in to comment.