diff --git a/.gitignore b/.gitignore
index d5c1656..817f8b2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -26,4 +26,6 @@ Network Trash Folder
Temporary Items
.apdisk
-.idea
\ No newline at end of file
+.idea
+
+.ipynb_checkpoints
\ No newline at end of file
diff --git a/README.md b/README.md
index a51e67d..909a4cb 100644
--- a/README.md
+++ b/README.md
@@ -4,11 +4,11 @@ Tool for visualizing BERT's attention layers
## Jupyter Notebooks
-[bertviz_detail](https://github.com/jessevig/bertviz/blob/master/bertviz/bertviz_detail.ipynb): new attention details view
+[bertviz_detail](https://github.com/jessevig/bertviz/blob/master/bertviz_detail.ipynb): new attention details view
* [Colab](https://colab.research.google.com/drive/1Nlhh2vwlQdKleNMqpmLDBsAwrv_7NnrB)
-[bertviz_summary](https://github.com/jessevig/bertviz/blob/master/bertviz/bertviz_summary.ipynb): original summary view
+[bertviz_summary](https://github.com/jessevig/bertviz/blob/master/bertviz_summary.ipynb): original summary view
* [Colab](https://colab.research.google.com/drive/1vlOJ1lhdujVjfH857hvYKIdKPTD9Kid8)
diff --git a/bertviz/bertviz_detail.ipynb b/bertviz/bertviz_detail.ipynb
deleted file mode 100644
index 43c0b97..0000000
--- a/bertviz/bertviz_detail.ipynb
+++ /dev/null
@@ -1,90 +0,0 @@
-{
- "cells": [
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "collapsed": true
- },
- "outputs": [],
- "source": [
- "from bertviz.attention_details import AttentionDetailsData, show\n",
- "from bertviz.pytorch_pretrained_bert import BertModel, BertTokenizer "
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "%%javascript\n",
- "require.config({\n",
- " paths: {\n",
- " d3: '//cdnjs.cloudflare.com/ajax/libs/d3/5.7.0/d3.min'\n",
- " }\n",
- "});"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "def call_html():\n",
- " import IPython\n",
- " display(IPython.core.display.HTML('''\n",
- " \n",
- " \n",
- " '''))"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "bert_version = 'bert-base-uncased'\n",
- "model = BertModel.from_pretrained(bert_version)\n",
- "tokenizer = BertTokenizer.from_pretrained(bert_version)\n",
- "sentence_a = \"I went to the store.\"\n",
- "sentence_b = \"At the store, I bought fresh strawberries.\"\n",
- "details_data = AttentionDetailsData(model, tokenizer)\n",
- "tokens_a, tokens_b, queries, keys, atts = details_data.get_data(sentence_a, sentence_b)\n",
- "call_html()\n",
- "show(tokens_a, tokens_b, queries, keys, atts)"
- ]
- }
- ],
- "metadata": {
- "kernelspec": {
- "display_name": "Python 2",
- "language": "python",
- "name": "python2"
- },
- "language_info": {
- "codemirror_mode": {
- "name": "ipython",
- "version": 2
- },
- "file_extension": ".py",
- "mimetype": "text/x-python",
- "name": "python",
- "nbconvert_exporter": "python",
- "pygments_lexer": "ipython2",
- "version": "2.7.6"
- }
- },
- "nbformat": 4,
- "nbformat_minor": 0
-}
diff --git a/bertviz/bertviz_summary.ipynb b/bertviz/bertviz_summary.ipynb
deleted file mode 100644
index 7b249f4..0000000
--- a/bertviz/bertviz_summary.ipynb
+++ /dev/null
@@ -1,80 +0,0 @@
-{
- "cells": [
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {
- "collapsed": true
- },
- "outputs": [],
- "source": [
- "from bertviz import attention, visualization\n",
- "from bertviz.pytorch_pretrained_bert import BertModel, BertTokenizer"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "\n",
- "%%javascript\n",
- "require.config({\n",
- " paths: {\n",
- " d3: '//cdnjs.cloudflare.com/ajax/libs/d3/3.4.8/d3.min'\n",
- " }\n",
- "});\n"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "def call_html():\n",
- " import IPython\n",
- " display(IPython.core.display.HTML('''\n",
- " \n",
- " \n",
- " '''))"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": [
- "bert_version = 'bert-base-uncased'\n",
- "model = BertModel.from_pretrained(bert_version)\n",
- "tokenizer = BertTokenizer.from_pretrained(bert_version)\n",
- "sentence_a = \"I went to the store.\"\n",
- "sentence_b = \"At the store, I bought fresh strawberries.\"\n",
- "attention_visualizer = visualization.AttentionVisualizer(model, tokenizer)\n",
- "tokens_a, tokens_b, attn = attention_visualizer.get_viz_data(sentence_a, sentence_b)\n",
- "call_html()\n",
- "attention.show(tokens_a, tokens_b, attn)"
- ]
- },
- {
- "cell_type": "code",
- "execution_count": null,
- "metadata": {},
- "outputs": [],
- "source": []
- }
- ],
- "metadata": {},
- "nbformat": 4,
- "nbformat_minor": 0
-}
diff --git a/bertviz_detail.ipynb b/bertviz_detail.ipynb
new file mode 100644
index 0000000..926a9bf
--- /dev/null
+++ b/bertviz_detail.ipynb
@@ -0,0 +1,1117 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "execution_count": 1,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from bertviz.attention_details import AttentionDetailsData, show\n",
+ "from bertviz.pytorch_pretrained_bert import BertModel, BertTokenizer "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 2,
+ "metadata": {},
+ "outputs": [
+ {
+ "data": {
+ "application/javascript": [
+ "require.config({\n",
+ " paths: {\n",
+ " d3: '//cdnjs.cloudflare.com/ajax/libs/d3/5.7.0/d3.min'\n",
+ " }\n",
+ "});"
+ ],
+ "text/plain": [
+ ""
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ }
+ ],
+ "source": [
+ "%%javascript\n",
+ "require.config({\n",
+ " paths: {\n",
+ " d3: '//cdnjs.cloudflare.com/ajax/libs/d3/5.7.0/d3.min'\n",
+ " }\n",
+ "});"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 3,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "def call_html():\n",
+ " import IPython\n",
+ " display(IPython.core.display.HTML('''\n",
+ " \n",
+ " \n",
+ " '''))"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": 4,
+ "metadata": {},
+ "outputs": [
+ {
+ "name": "stderr",
+ "output_type": "stream",
+ "text": [
+ "02/03/2019 08:42:44 - INFO - bertviz.pytorch_pretrained_bert.modeling - loading archive file https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased.tar.gz from cache at /Users/jvig/.pytorch_pretrained_bert/9c41111e2de84547a463fd39217199738d1e3deb72d4fec4399e6e241983c6f0.ae3cef932725ca7a30cdcb93fc6e09150a55e2a130ec7af63975a16c153ae2ba\n",
+ "02/03/2019 08:42:44 - INFO - bertviz.pytorch_pretrained_bert.modeling - extracting archive file /Users/jvig/.pytorch_pretrained_bert/9c41111e2de84547a463fd39217199738d1e3deb72d4fec4399e6e241983c6f0.ae3cef932725ca7a30cdcb93fc6e09150a55e2a130ec7af63975a16c153ae2ba to temp dir /var/folders/dn/mjpk8sk945q2m1c1lppnq4ym0000gn/T/tmpyr_5vpeg\n",
+ "02/03/2019 08:42:48 - INFO - bertviz.pytorch_pretrained_bert.modeling - Model config {\n",
+ " \"attention_probs_dropout_prob\": 0.1,\n",
+ " \"hidden_act\": \"gelu\",\n",
+ " \"hidden_dropout_prob\": 0.1,\n",
+ " \"hidden_size\": 768,\n",
+ " \"initializer_range\": 0.02,\n",
+ " \"intermediate_size\": 3072,\n",
+ " \"max_position_embeddings\": 512,\n",
+ " \"num_attention_heads\": 12,\n",
+ " \"num_hidden_layers\": 12,\n",
+ " \"type_vocab_size\": 2,\n",
+ " \"vocab_size\": 30522\n",
+ "}\n",
+ "\n",
+ "02/03/2019 08:42:50 - INFO - bertviz.pytorch_pretrained_bert.tokenization - loading vocabulary file https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-vocab.txt from cache at /Users/jvig/.pytorch_pretrained_bert/26bc1ad6c0ac742e9b52263248f6d0f00068293b33709fae12320c0e35ccfbbb.542ce4285a40d23a559526243235df47c5f75c197f04f37d1a0c124c32c9a084\n"
+ ]
+ },
+ {
+ "data": {
+ "text/html": [
+ "\n",
+ " \n",
+ " \n",
+ " "
+ ],
+ "text/plain": [
+ ""
+ ]
+ },
+ "metadata": {},
+ "output_type": "display_data"
+ },
+ {
+ "data": {
+ "text/html": [
+ "\n",
+ " \n",
+ " Layer: \n",
+ " Head: \n",
+ " Attention: \n",
+ " \n",
+ " \n"
+ ],
+ "text/plain": [
+ "