Use PyTorch UAE in imdb example (#705)

SeldonIO · Jan 11, 2023 · 97a9f41 · 97a9f41
1 parent 3daf62c
commit 97a9f41
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/doc/source/examples/cd_text_imdb.ipynb b/doc/source/examples/cd_text_imdb.ipynb
@@ -974,15 +974,14 @@
    "source": [
     "from alibi_detect.cd.pytorch import preprocess_drift\n",
     "from alibi_detect.models.pytorch import TransformerEmbedding\n",
+    "from alibi_detect.cd.pytorch import UAE\n",
     "\n",
+    "# Embedding model\n",
     "embedding_pt = TransformerEmbedding(model_name, emb_type, layers)\n",
     "\n",
-    "model = nn.Sequential(\n",
-    "    embedding_pt,\n",
-    "    nn.Linear(768, 256),\n",
-    "    nn.ReLU(),\n",
-    "    nn.Linear(256, enc_dim)\n",
-    ").to(device).eval()\n",
+    "# PyTorch untrained autoencoder\n",
+    "uae = UAE(input_layer=embedding_pt, shape=shape, enc_dim=enc_dim)\n",
+    "model = uae.to(device).eval()\n",
     "\n",
     "# define preprocessing function\n",
     "preprocess_fn = partial(preprocess_drift, model=model, tokenizer=tokenizer, \n",