From 5b8902ca1f4499fa6f3e156af9eb01b9982cbb58 Mon Sep 17 00:00:00 2001
From: YannDubs <yanndubois96@gmail.com>
Date: Mon, 5 Jul 2021 14:12:07 -0400
Subject: [PATCH] [HUB] add docstring.

---
 hub/__init__.py   |  2 +-
 hub/compressor.py |  1 +
 hubconf.py        | 18 +++++++++++++-----
 3 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/hub/__init__.py b/hub/__init__.py
index 0f8fc49..1a7ca40 100644
--- a/hub/__init__.py
+++ b/hub/__init__.py
@@ -1 +1 @@
-from compressor import *
+from .compressor import *
diff --git a/hub/compressor.py b/hub/compressor.py
index 662ab12..e9ccd55 100644
--- a/hub/compressor.py
+++ b/hub/compressor.py
@@ -64,6 +64,7 @@ def __init__(
 
         self.device = device
         self.to(self.device)
+        self.eval()
 
     def to(self, device):
         self.device = device
diff --git a/hubconf.py b/hubconf.py
index 8dfd48e..b493676 100644
--- a/hubconf.py
+++ b/hubconf.py
@@ -2,6 +2,8 @@
     "torch",
     "compressai",
     "clip",
+    "tqdm",
+    "numpy",
 ]  # dependencies required for loading a model
 
 import os
@@ -12,7 +14,7 @@
 
 
 def clip_compressor(beta=5e-02, pretrained=True, **kwargs):
-    """ 
+    """Load an invariant CLIP compressor.
 
     Parameters
     ----------
@@ -31,10 +33,16 @@ def clip_compressor(beta=5e-02, pretrained=True, **kwargs):
 
     Return
     ------
-    compressor : 
-
-    transform : 
-
+    compressor : nn.Module
+        Pytorch module that when called as `compressor(X)` on a batch of image, will return
+        decompressed representations. Use `compressor.compress(X)` to get a batch of compressed 
+        representations (in bytes). To save a compressed torch dataset to file use 
+        `compressor.compress_dataset(dataset,file)` and `dataset = compressor.decompress_dataset(file)`.
+        For more information check the docstrings of all functions of the module and the examples below.
+
+    transform : callable
+        Transforms that can be used directly in place of torchvision transform. It will resize the
+        image to (3,224,224), apply clip normalization and convert it to tensor.
     """
     base = "https://github.com/YannDubs/lossyless/releases/download"
     ckpt = f"{base}/v0.1-alpha/beta{beta:0.0e}_factorized_rate.pt"