Fix TF 2x LLM SQ Legacy Keras Environment Variable Issue (#1818)

zehao-intel · web-flow · commit 27644940e746 · 2024-05-28T13:38:27.000+08:00
Signed-off-by: zehao-intel &lt;zehao.huang@intel.com&gt;
diff --git a/examples/tensorflow/nlp/large_language_models/quantization/ptq/smoothquant/benchmark.py b/examples/tensorflow/nlp/large_language_models/quantization/ptq/smoothquant/benchmark.py
@@ -167,18 +167,9 @@ def __len__(self):
 else:
     print("benchmaking fp32 model")
     model = transformers.TFAutoModelForCausalLM.from_pretrained(model_name)
-    # fp32_folder = model_name.split('/')[-1] + "_fp32"
-    # model.save(fp32_folder)
-    # model = tf.keras.models.load_model(fp32_folder)
     from neural_compressor.experimental import common
-    def keras2SavedModel(model):
-        model = common.Model(model)
-        return model.model
-    model = keras2SavedModel(model) # tensorflow.python.trackable.autotrackable.AutoTrackable object
 
-# TODO current neural_compressor.benchmark does not support AutoTrackable model, we will write our own
-# from neural_compressor.benchmark import fit
-# from neural_compressor.config import BenchmarkConfig
-# conf = BenchmarkConfig(cores_per_instance=28, num_of_instance=1)
-# fit(model, conf, b_func=evaluator.evaluate_tf_v1)
+    os.environ["TF_USE_LEGACY_KERAS"]="False"
+    model = common.Model(model).model # tensorflow.python.trackable.autotrackable.AutoTrackable object
+
 evaluator.evaluate_tf_v1(model)
diff --git a/examples/tensorflow/nlp/large_language_models/quantization/ptq/smoothquant/main.py b/examples/tensorflow/nlp/large_language_models/quantization/ptq/smoothquant/main.py
@@ -15,7 +15,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-import os.path
+import os
 import transformers
 import tensorflow as tf
 from tqdm import tqdm
@@ -186,9 +186,9 @@ def eval_func(model):
 
 from neural_compressor import PostTrainingQuantConfig
 from neural_compressor.config import AccuracyCriterion
-
 from neural_compressor import quantization
 
+os.environ["TF_USE_LEGACY_KERAS"]="False"
 recipes = {}
 if args.sq:
     recipes = {"smooth_quant": True, "smooth_quant_args": {'alpha': args.alpha}}
