Update PyPI

Author: yongchanghao

MANIFEST.in

@@ -0,0 +1,2 @@
+recursive-include cuda *.cu 
+recursive-include cuda *.cuh 

README.md

@@ -1,29 +1,41 @@
 # NeuZip: Memory-Efficient Training and Inference with Dynamic Compression of Neural Networks
 
-This is the official repository for the paper "NeuZip: Memory-Efficient Training and Inference with Dynamic Compression of Neural Networks".
-This repository contains the code for the experiments in the paper.
+This is the official repository for the paper [*NeuZip: Memory-Efficient Training and Inference with Dynamic Compression of Neural Networks*](https://arxiv.org/abs/2410.20650). This repository contains the code for the experiments in the paper.
 
 # Installation
 
-First, please install NVComp library on your own. You can find the installation instructions on its [GitHub page](https://github.com/NVIDIA/nvcomp).
+## From PyPI
 
-Then, you can install the package by running the following command in the root directory of the repository:
+We provide a simple way to install the package from PyPI. You can install the package by running the following command:
 
 ```bash
+pip install neuzip
+```
+
+Note that we only upload the source distribution to PyPI. You need to have NVCC correctly installed on your system to compile the package.
+
+## From source
+
+You can also install the package from source, which is useful if you want to modify the code.
+
+```bash
+git clone https://github.com/BorealisAI/neuzip
+cd neuzip
 pip install -e .
 ```
 
-# Usage
+# Basic usage
 
-You can use the `neuzip` package to compress and decompress tensors. Here is an example:
+Using `neuzip` for your PyTorch model is pretty easy. Here is a simple example:
 
-```python
-import neuzip
-manager = neuzip.Manager()  # Create a manager
-handle = manager.write(tensor)  # Compress a tensor
-tensor = manager.read(handle)  # Decompress a tensor
+```diff
+model: torch.nn.Module = # your model
++ manager = neuzip.Manager()
++ model = manager.convert(model)
 ```
 
-# Replicating Experiments
+The compressed model can be used in the same way as the original model while consuming less memory.
+
+# Replicating experiments
 
 You can replicate all the experiments in the paper by using the files in the [examples/](examples/) directory. Each file corresponds to one or more experiments in the paper.

pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools", "torch", "numpy"]
+build-backend = "setuptools.build_meta"

setup.py

@@ -8,10 +8,11 @@
 import os
 
 import torch
-from setuptools import setup
+from setuptools import find_packages, setup
+from torch.torch_version import Version
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 
-NEUZIP_VERSION = "0.0.1"
+NEUZIP_VERSION = "0.1.1"
 
 setup(
     name="neuzip",
@@ -33,6 +34,19 @@
             ],
         ),
     ],
+    install_requires=[
+        "torch",
+        "transformers",
+        "numpy",
+        f"nvidia-nvcomp-cu{Version(torch.version.cuda).major}",
+    ],
+    packages=find_packages(),
+    package_data={
+        "neuzip._cuda": ["*.cuh", "*.cu"],
+    },
+    include_package_data=True,
+    url="https://github.com/BorealisAI/neuzip",
+    license="CC BY-NC-SA 4.0",
     version=NEUZIP_VERSION,
     cmdclass={"build_ext": BuildExtension},
 )