This repository is a benchmark for 'PUMA' based on secret flow. To run it, you have to install the following packages:
jax==0.4.28
flax==0.8.3
secretflow==1.5.0b0
spu==0.9.0b1
Maybe run pip install secretflow
is enough.
And also transformers[flax]
shall be installed.