Merge pull request #191 from stochasticai/dev

dev

Author: StochasticRomanAgeev

src/xturing/datasets/instruction_dataset.py

@@ -59,7 +59,6 @@ def __init__(
         else:
             path = Path(path)
             assert Path(path).exists(), "path does not exist"
-
             if path.is_dir():
                 self.data = load_from_disk(str(path))
             elif path.suffix == ".jsonl":
@@ -123,7 +122,7 @@ def __getitem__(self, idx):
         return self.data["train"][idx]
 
     def save(self, path):
-        return self.data.save_to_disk(path)
+        return self.data["train"].save_to_disk(path)
 
     @classmethod
     def generate_dataset(

src/xturing/engines/causal.py

@@ -175,6 +175,9 @@ def __init__(
 
         self.loss_fct = CrossEntropyLoss()
 
+    def set_from_state_dict(self, state_dict, strict=False):
+        self.model.load_state_dict(state_dict, strict=strict)
+
     def save(self, saving_path: Union[str, Path]):
         # Save HF config file
         self.model.config.save_pretrained(str(saving_path))

src/xturing/engines/llama_engine.py

@@ -161,10 +161,6 @@ def noop(*args, **kwargs):
         state_dict = torch.load(
             weights_path / Path("pytorch_model.bin"), map_location="cpu"
         )
-        new_state_dict = {}
-        for key, value in state_dict.items():
-            new_state_dict[key[6:]] = value
-        model.load_state_dict(new_state_dict, strict=False)
 
         if warmup_autotune:
             autotune_warmup(model)
@@ -192,3 +188,5 @@ def noop(*args, **kwargs):
         torch.nn.init.kaiming_uniform_ = saved_kaiming_uniform_
         torch.nn.init.uniform_ = saved_uniform_
         torch.nn.init.normal_ = saved_normal_
+
+        self.set_from_state_dict(state_dict)

src/xturing/utils/hub.py

@@ -44,7 +44,7 @@ def bar_progress(current, total, width=80):
                 sys.stdout.flush()
 
             try:
-                wget.download(url, str(zip_filename), bar=bar_progress)
+                wget.download(url, str(zip_filename))
 
                 with ZipFile(zip_filename, "r") as zip_ref:
                     zip_ref.extractall(path=model_dir)

tests/xturing/datasets/test_instruction_dataset.py

@@ -42,3 +42,4 @@ def test_features_dataset():
         "target": "second text",
         "instruction": "second instruction",
     }
+    dataset.save(".")