fix(decoder): remove call to beta serialization api (#393)

* fix(decoder): remove call to beta serialization api * fix: style issue in notebooks
huggingface · Dec 22, 2023 · 63aa3b7 · 63aa3b7
1 parent cd59cac
commit 63aa3b7
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 7 deletions.
diff --git a/notebooks/text-generation/scripts/run_clm.py b/notebooks/text-generation/scripts/run_clm.py
@@ -1,14 +1,13 @@
 from dataclasses import dataclass, field
-import os
+
+from datasets import load_from_disk
 from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
-    set_seed,
     default_data_collator,
+    set_seed,
 )
-from datasets import load_from_disk
-import torch
-from huggingface_hub import login
+
 from optimum.neuron import NeuronHfArgumentParser as HfArgumentParser
 from optimum.neuron import NeuronTrainer as Trainer
 from optimum.neuron import NeuronTrainingArguments as TrainingArguments

diff --git a/notebooks/text-generation/scripts/utils/pack_dataset.py b/notebooks/text-generation/scripts/utils/pack_dataset.py
@@ -1,5 +1,6 @@
-from itertools import chain
 from functools import partial
+from itertools import chain
+
 
 remainder = {"input_ids": [], "attention_mask": [], "token_type_ids": []}
 

diff --git a/optimum/neuron/modeling_decoder.py b/optimum/neuron/modeling_decoder.py
@@ -246,7 +246,7 @@ def _save_pretrained(self, save_directory: Union[str, Path]):
 
         if src_compiled_path is None:
             # The compiled model has never been serialized: do it now
-            self.model._save_compiled_artifacts(dst_compiled_path)
+            self.model.save(dst_compiled_path)
         else:
             shutil.copytree(src_compiled_path, dst_compiled_path)