[Safetensors] Fix fast safe open slice. (#8512)

* fix more case.

* fix slice.

* fix error slice.

Author: ZHUI

paddlenlp/utils/safetensors.py

@@ -157,16 +157,16 @@ def __getitem__(self, index):
 
         out_start, out_stop, out_step = copy.deepcopy((self.start, self.stop, self.step))
         for i, (start, stop, step, slice_) in enumerate(zip(self.start, self.stop, self.step, index)):
-            out_start[i] = slice_.start or 0
-            out_step[i] = slice_.step or 1
-            out_stop[i] = slice_.stop or stop - start
+            out_start[i] = slice_.start if slice_.start is not None else 0
+            out_step[i] = slice_.step if slice_.step is not None else 1
+            out_stop[i] = slice_.stop if slice_.stop is not None else stop - start
             out_stop[i] = min(stop, out_stop[i])
 
         target_shape = []
-        for x, y, z in zip(out_start, out_stop, out_step):
+        for x, y, z, sli in zip(out_start, out_stop, out_step, index):
             assert z == 1, "only support step = 1"
-            if y - x > 1:
-                target_shape.append(int(y - x))
+            if y - x > 1 or sli.step is None:
+                target_shape.append(max(int(y - x), 0))
 
         if len(target_shape) == 0:
             if self.shape == [1]:

tests/transformers/test_safetensors.py

@@ -28,7 +28,14 @@ class FastSafetensors(unittest.TestCase):
     def setUp(self):
         super().setUp()
         self.weigth_map = {}
-        tensors = [([10, 10], "float32"), ([8], "float16"), ([5, 5, 5], "int32")]
+        tensors = [
+            ([10, 1, 10], "float32"),
+            ([1, 1, 10], "float32"),
+            ([1, 1, 1, 10], "float32"),
+            ([10, 10], "float32"),
+            ([8], "float16"),
+            ([5, 5, 5], "int32"),
+        ]
         count = 0
         for shape, dtype in tensors:
             self.weigth_map[f"weight_{count}"] = (np.random.random(shape) * 100).astype(dtype)
@@ -53,5 +60,10 @@ def test_safe_open(self):
             with fast_safe_open(path, framework="np") as f:
                 for key in f.keys():
                     safe_slice = f.get_slice(key)
+                    # np.testing.assert_equal(self.weigth_map[key][2:1, ...], safe_slice[2:1, ...])
+                    np.testing.assert_equal(self.weigth_map[key][0, ...], safe_slice[0, ...])
+                    np.testing.assert_equal(self.weigth_map[key][0:1, ...], safe_slice[0:1, ...])
+                    np.testing.assert_equal(self.weigth_map[key][..., 2:], safe_slice[..., 2:])
+                    np.testing.assert_equal(self.weigth_map[key][..., 1], safe_slice[..., 1])
                     np.testing.assert_equal(self.weigth_map[key][:2, ...], safe_slice[:2, ...])
                     np.testing.assert_equal(self.weigth_map[key][..., :4], safe_slice[..., :4])