- Include transformers in requirements

- run black
- fix imports

examples/HuggingFace_ResNet_Implementation.ipynb

@@ -27,7 +27,7 @@
     "import optax\n",
     "\n",
     "from znnl.models import HuggingFaceFlaxModel\n",
-    "\n",
+    "from transformers import ResNetConfig, FlaxResNetForImageClassification\n",
     "import jax\n",
     "print(jax.default_backend())"
    ]
@@ -174,6 +174,20 @@
     "plt.legend()\n",
     "plt.show()"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

requirements.txt

@@ -13,7 +13,7 @@ tensorflow_probability
 scipy
 scikit-learn
 jaxlib
-jax
+# jax
 plotly
 flax
 tqdm
@@ -23,4 +23,5 @@ tensorflow-datasets
 isort
 tensorflow
 pyyaml
-jupyter
+jupyter
+transformers

znnl/models/__init__.py

@@ -24,6 +24,7 @@
 Summary
 -------
 """
+
 from znnl.models.flax_model import FlaxModel
 from znnl.models.huggingface_flax_model import HuggingFaceFlaxModel
 from znnl.models.jax_model import JaxModel

znnl/models/flax_model.py

@@ -24,6 +24,7 @@
 Summary
 -------
 """
+
 import logging
 from typing import Callable, List, Sequence, Union
 

znnl/models/huggingface_flax_model.py

@@ -25,13 +25,14 @@
 -------
 Module for using a Flax model from Hugging Face in ZnNL.
 """
+
 import logging
 from typing import Callable, List, Sequence, Union
 
 import jax
 import jax.numpy as np
 from flax import linen as nn
-from transformers import FlaxPreTrainedModel, ResNetConfig
+from transformers import FlaxPreTrainedModel
 
 from znnl.models.jax_model import JaxModel
 

znnl/models/jax_model.py

@@ -24,6 +24,7 @@
 Summary
 -------
 """
+
 from typing import TYPE_CHECKING, Callable, Optional, Sequence, Union
 
 import jax
@@ -77,16 +78,16 @@ def __init__(
         self.optimizer = optimizer
         self.input_shape = input_shape
 
+        # Initialized in self.init_model
+        self.rng = None
+
         # Input shape is required if no full model is passed.
         if pre_built_model is None and input_shape is None:
             raise ValueError(
                 "Input shape must be specified if no pre-built model is passed."
                 "Model is yet to be constructed."
             )
 
-        # Initialized in self.init_model
-        self.rng = None
-
         # initialize the model state
         if pre_built_model is None:
             self.init_model(seed=seed)

znnl/models/nt_model.py

@@ -24,6 +24,7 @@
 Summary
 -------
 """
+
 import logging
 from typing import Callable, Sequence, Union