repo-sync-2023-12-12T14:27:43+0800

anakinxc · Dec 12, 2023 · 05cafcd · 05cafcd
1 parent 9f6d31b
commit 05cafcd
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 6 deletions.
diff --git a/bazel/repositories.bzl b/bazel/repositories.bzl
@@ -20,7 +20,7 @@ SECRETFLOW_GIT = "https://github.com/secretflow"
 
 YACL_COMMIT_ID = "6ba8bd5f02035176ec4daaca1c1269195a1b1b4e"
 
-LIBSPI_COMMIT_ID = "575cb5999b5bc52926d08de6944e8d54f0a38f9f"
+LIBSPI_COMMIT_ID = "dbf452b9d87619b40b73746b7c9afefc1845975b"
 
 def spu_deps():
     _rules_cuda()

diff --git a/examples/python/ml/README.md b/examples/python/ml/README.md
@@ -13,7 +13,8 @@ This directory contains examples demonstrating how to use SPU to write privacy-p
                        [Stax](https://jax.readthedocs.io/en/latest/jax.example_libraries.stax.html).
 * [stax_mnist_classifier](stax_mnist_classifier/): Private training of a simple neural network for MNIST classification with
                                                    [Stax](https://jax.readthedocs.io/en/latest/jax.example_libraries.stax.html).
-* [flax_resnet](flax_resnet/): Private training of a [ResNet](https://arxiv.org/abs/1512.03385) model with [Flax](https://github.com/google/flax) library, and private inference of a pre-trained ResNet-50 model based on [Microsoft-ResNet50](https://huggingface.co/microsoft/resnet-50).
+* [flax_resnet](flax_resnet/): Private training of a [ResNet](https://arxiv.org/abs/1512.03385) model with [Flax](https://github.com/google/flax)
+library, and private inference of a pre-trained ResNet-50 model based on [Microsoft-ResNet50](https://huggingface.co/microsoft/resnet-50).
 * [flax_gpt2](flax_gpt2/): Private inference of a pre-trained
                            [GPT2](https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf)
                            model with [Flax](https://github.com/google/flax) library.

diff --git a/spu/pir.py b/spu/pir.py
@@ -16,7 +16,7 @@
 
 from typing import List
 
-from . import libspu  # type: ignore
+from .libspu.link import Context  # type: ignore
 from . import libpsi  # type: ignore
 from .pir_pb2 import (  # type: ignore
     KvStoreType,

diff --git a/spu/psi.py b/spu/psi.py
@@ -16,7 +16,8 @@
 
 from typing import List
 
-from . import libspu  # type: ignore
+from .libspu.link import Context  # type: ignore
+from .libpsi.libs import ProgressData
 from . import libpsi  # type: ignore
 from .psi_pb2 import (  # type: ignore
     BucketPsiConfig,
@@ -30,13 +31,13 @@
 
 
 def mem_psi(
-    link: libspu.link.Context, config: MemoryPsiConfig, input_items: List[str]
+    link: Context, config: MemoryPsiConfig, input_items: List[str]
 ) -> List[str]:
     return libpsi.libs.mem_psi(link, config.SerializeToString(), input_items)
 
 
 def bucket_psi(
-    link: libspu.link.Context,
+    link: Context,
     config: BucketPsiConfig,
     progress_callbacks: [[libpsi.libs.ProgressData], None] = None,
     callbacks_interval_ms: int = 5 * 1000,