draw.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/svg+xml": [
       "<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n",
       "<!DOCTYPE svg PUBLIC \"-//W3C//DTD SVG 1.1//EN\"\n",
       " \"http://www.w3.org/Graphics/SVG/1.1/DTD/svg11.dtd\">\n",
       "<!-- Generated by graphviz version 12.0.0 (20240704.0754)\n",
       " -->\n",
       "<!-- Title: model Pages: 1 -->\n",
       "<svg width=\"503pt\" height=\"353pt\"\n",
       " viewBox=\"0.00 0.00 502.50 352.50\" xmlns=\"http://www.w3.org/2000/svg\" xmlns:xlink=\"http://www.w3.org/1999/xlink\">\n",
       "<g id=\"graph0\" class=\"graph\" transform=\"scale(1 1) rotate(0) translate(4 348.5)\">\n",
       "<title>model</title>\n",
       "<polygon fill=\"white\" stroke=\"none\" points=\"-4,4 -4,-348.5 498.5,-348.5 498.5,4 -4,4\"/>\n",
       "<!-- 0 -->\n",
       "<g id=\"node1\" class=\"node\">\n",
       "<title>0</title>\n",
       "<polygon fill=\"lightyellow\" stroke=\"none\" points=\"215.38,-339.75 116.62,-339.75 116.62,-307.25 215.38,-307.25 215.38,-339.75\"/>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"116.62,-307.25 116.62,-339.75 177.62,-339.75 177.62,-307.25 116.62,-307.25\"/>\n",
       "<text text-anchor=\"start\" x=\"121.62\" y=\"-325.25\" font-family=\"Linux libertine\" font-size=\"10.00\">input&#45;tensor</text>\n",
       "<text text-anchor=\"start\" x=\"131\" y=\"-314\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:0</text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"177.62,-307.25 177.62,-339.75 215.38,-339.75 215.38,-307.25 177.62,-307.25\"/>\n",
       "<text text-anchor=\"start\" x=\"182.62\" y=\"-319.62\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10)</text>\n",
       "</g>\n",
       "<!-- 1 -->\n",
       "<g id=\"node2\" class=\"node\">\n",
       "<title>1</title>\n",
       "<polygon fill=\"aliceblue\" stroke=\"none\" points=\"146,-266.5 0,-266.5 0,-224.5 146,-224.5 146,-266.5\"/>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"0,-224.5 0,-266.5 63,-266.5 63,-224.5 0,-224.5\"/>\n",
       "<text text-anchor=\"start\" x=\"4.88\" y=\"-247.25\" font-family=\"Linux libertine\" font-size=\"10.00\">__getitem__</text>\n",
       "<text text-anchor=\"start\" x=\"15.38\" y=\"-236\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"63,-245.5 63,-266.5 106,-266.5 106,-245.5 63,-245.5\"/>\n",
       "<text text-anchor=\"start\" x=\"72.5\" y=\"-252\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"106,-245.5 106,-266.5 146,-266.5 146,-245.5 106,-245.5\"/>\n",
       "<text text-anchor=\"start\" x=\"110.62\" y=\"-252\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"63,-224.5 63,-245.5 106,-245.5 106,-224.5 63,-224.5\"/>\n",
       "<text text-anchor=\"start\" x=\"68\" y=\"-231\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"106,-224.5 106,-245.5 146,-245.5 146,-224.5 106,-224.5\"/>\n",
       "<text text-anchor=\"start\" x=\"113.25\" y=\"-231\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 2) </text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;1 -->\n",
       "<g id=\"edge1\" class=\"edge\">\n",
       "<title>0&#45;&gt;1</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M147.63,-307.48C135.81,-297.83 120.13,-285.01 106.18,-273.62\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"108.46,-270.96 98.5,-267.34 104.03,-276.38 108.46,-270.96\"/>\n",
       "</g>\n",
       "<!-- 2 -->\n",
       "<g id=\"node3\" class=\"node\">\n",
       "<title>2</title>\n",
       "<polygon fill=\"#c1ffc1\" stroke=\"none\" points=\"351.5,-266.5 164.5,-266.5 164.5,-224.5 351.5,-224.5 351.5,-266.5\"/>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"164.5,-224.5 164.5,-266.5 246.5,-266.5 246.5,-224.5 164.5,-224.5\"/>\n",
       "<text text-anchor=\"start\" x=\"169.5\" y=\"-247.25\" font-family=\"Linux libertine\" font-size=\"10.00\">BertEmbeddings</text>\n",
       "<text text-anchor=\"start\" x=\"189.38\" y=\"-236\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"246.5,-245.5 246.5,-266.5 289.5,-266.5 289.5,-245.5 246.5,-245.5\"/>\n",
       "<text text-anchor=\"start\" x=\"256\" y=\"-252\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"289.5,-245.5 289.5,-266.5 351.5,-266.5 351.5,-245.5 289.5,-245.5\"/>\n",
       "<text text-anchor=\"start\" x=\"305.12\" y=\"-252\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"246.5,-224.5 246.5,-245.5 289.5,-245.5 289.5,-224.5 246.5,-224.5\"/>\n",
       "<text text-anchor=\"start\" x=\"251.5\" y=\"-231\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"289.5,-224.5 289.5,-245.5 351.5,-245.5 351.5,-224.5 289.5,-224.5\"/>\n",
       "<text text-anchor=\"start\" x=\"294.25\" y=\"-231\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10, 384) </text>\n",
       "</g>\n",
       "<!-- 0&#45;&gt;2 -->\n",
       "<g id=\"edge2\" class=\"edge\">\n",
       "<title>0&#45;&gt;2</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M184.18,-307.48C195.86,-297.83 211.38,-285.01 225.18,-273.62\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"227.29,-276.41 232.77,-267.35 222.83,-271.02 227.29,-276.41\"/>\n",
       "</g>\n",
       "<!-- 5 -->\n",
       "<g id=\"node6\" class=\"node\">\n",
       "<title>5</title>\n",
       "<polygon fill=\"#c1ffc1\" stroke=\"none\" points=\"342.5,-188.5 173.5,-188.5 173.5,-146.5 342.5,-146.5 342.5,-188.5\"/>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"173.5,-146.5 173.5,-188.5 237.5,-188.5 237.5,-146.5 173.5,-146.5\"/>\n",
       "<text text-anchor=\"start\" x=\"178.5\" y=\"-169.25\" font-family=\"Linux libertine\" font-size=\"10.00\">BertEncoder</text>\n",
       "<text text-anchor=\"start\" x=\"189.38\" y=\"-158\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"237.5,-167.5 237.5,-188.5 280.5,-188.5 280.5,-167.5 237.5,-167.5\"/>\n",
       "<text text-anchor=\"start\" x=\"247\" y=\"-174\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"280.5,-167.5 280.5,-188.5 342.5,-188.5 342.5,-167.5 280.5,-167.5\"/>\n",
       "<text text-anchor=\"start\" x=\"285.25\" y=\"-174\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10, 384) </text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"237.5,-146.5 237.5,-167.5 280.5,-167.5 280.5,-146.5 237.5,-146.5\"/>\n",
       "<text text-anchor=\"start\" x=\"242.5\" y=\"-153\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"280.5,-146.5 280.5,-167.5 342.5,-167.5 342.5,-146.5 280.5,-146.5\"/>\n",
       "<text text-anchor=\"start\" x=\"285.25\" y=\"-153\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10, 384) </text>\n",
       "</g>\n",
       "<!-- 2&#45;&gt;5 -->\n",
       "<g id=\"edge3\" class=\"edge\">\n",
       "<title>2&#45;&gt;5</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M258,-224.53C258,-216.94 258,-208.1 258,-199.71\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"261.5,-199.85 258,-189.85 254.5,-199.85 261.5,-199.85\"/>\n",
       "</g>\n",
       "<!-- 3 -->\n",
       "<g id=\"node4\" class=\"node\">\n",
       "<title>3</title>\n",
       "<polygon fill=\"aliceblue\" stroke=\"none\" points=\"494.5,-344.5 369.5,-344.5 369.5,-302.5 494.5,-302.5 494.5,-344.5\"/>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"369.5,-302.5 369.5,-344.5 411.5,-344.5 411.5,-302.5 369.5,-302.5\"/>\n",
       "<text text-anchor=\"start\" x=\"374.75\" y=\"-325.25\" font-family=\"Linux libertine\" font-size=\"10.00\">__eq__</text>\n",
       "<text text-anchor=\"start\" x=\"374.38\" y=\"-314\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"411.5,-323.5 411.5,-344.5 454.5,-344.5 454.5,-323.5 411.5,-323.5\"/>\n",
       "<text text-anchor=\"start\" x=\"421\" y=\"-330\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"454.5,-323.5 454.5,-344.5 494.5,-344.5 494.5,-323.5 454.5,-323.5\"/>\n",
       "<text text-anchor=\"start\" x=\"459.12\" y=\"-330\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"411.5,-302.5 411.5,-323.5 454.5,-323.5 454.5,-302.5 411.5,-302.5\"/>\n",
       "<text text-anchor=\"start\" x=\"416.5\" y=\"-309\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"454.5,-302.5 454.5,-323.5 494.5,-323.5 494.5,-302.5 454.5,-302.5\"/>\n",
       "<text text-anchor=\"start\" x=\"459.12\" y=\"-309\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
       "</g>\n",
       "<!-- 4 -->\n",
       "<g id=\"node5\" class=\"node\">\n",
       "<title>4</title>\n",
       "<polygon fill=\"aliceblue\" stroke=\"none\" points=\"494.5,-266.5 369.5,-266.5 369.5,-224.5 494.5,-224.5 494.5,-266.5\"/>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"369.5,-224.5 369.5,-266.5 411.5,-266.5 411.5,-224.5 369.5,-224.5\"/>\n",
       "<text text-anchor=\"start\" x=\"385.62\" y=\"-247.25\" font-family=\"Linux libertine\" font-size=\"10.00\">all</text>\n",
       "<text text-anchor=\"start\" x=\"374.38\" y=\"-236\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"411.5,-245.5 411.5,-266.5 454.5,-266.5 454.5,-245.5 411.5,-245.5\"/>\n",
       "<text text-anchor=\"start\" x=\"421\" y=\"-252\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"454.5,-245.5 454.5,-266.5 494.5,-266.5 494.5,-245.5 454.5,-245.5\"/>\n",
       "<text text-anchor=\"start\" x=\"459.12\" y=\"-252\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10) </text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"411.5,-224.5 411.5,-245.5 454.5,-245.5 454.5,-224.5 411.5,-224.5\"/>\n",
       "<text text-anchor=\"start\" x=\"416.5\" y=\"-231\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"454.5,-224.5 454.5,-245.5 494.5,-245.5 494.5,-224.5 454.5,-224.5\"/>\n",
       "<text text-anchor=\"start\" x=\"470\" y=\"-231\" font-family=\"Linux libertine\" font-size=\"10.00\">() </text>\n",
       "</g>\n",
       "<!-- 3&#45;&gt;4 -->\n",
       "<g id=\"edge4\" class=\"edge\">\n",
       "<title>3&#45;&gt;4</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M432,-302.53C432,-294.94 432,-286.1 432,-277.71\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"435.5,-277.85 432,-267.85 428.5,-277.85 435.5,-277.85\"/>\n",
       "</g>\n",
       "<!-- 6 -->\n",
       "<g id=\"node7\" class=\"node\">\n",
       "<title>6</title>\n",
       "<polygon fill=\"#c1ffc1\" stroke=\"none\" points=\"258,-110.5 98,-110.5 98,-68.5 258,-68.5 258,-110.5\"/>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"98,-68.5 98,-110.5 153,-110.5 153,-68.5 98,-68.5\"/>\n",
       "<text text-anchor=\"start\" x=\"102.62\" y=\"-91.25\" font-family=\"Linux libertine\" font-size=\"10.00\">BertPooler</text>\n",
       "<text text-anchor=\"start\" x=\"109.38\" y=\"-80\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:1</text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"153,-89.5 153,-110.5 196,-110.5 196,-89.5 153,-89.5\"/>\n",
       "<text text-anchor=\"start\" x=\"162.5\" y=\"-96\" font-family=\"Linux libertine\" font-size=\"10.00\">input:</text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"196,-89.5 196,-110.5 258,-110.5 258,-89.5 196,-89.5\"/>\n",
       "<text text-anchor=\"start\" x=\"200.75\" y=\"-96\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10, 384) </text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"153,-68.5 153,-89.5 196,-89.5 196,-68.5 153,-68.5\"/>\n",
       "<text text-anchor=\"start\" x=\"158\" y=\"-75\" font-family=\"Linux libertine\" font-size=\"10.00\">output: </text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"196,-68.5 196,-89.5 258,-89.5 258,-68.5 196,-68.5\"/>\n",
       "<text text-anchor=\"start\" x=\"209\" y=\"-75\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 384) </text>\n",
       "</g>\n",
       "<!-- 5&#45;&gt;6 -->\n",
       "<g id=\"edge5\" class=\"edge\">\n",
       "<title>5&#45;&gt;6</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M236.98,-146.53C227.79,-137.8 216.87,-127.43 206.93,-117.99\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"209.64,-115.73 199.98,-111.38 204.82,-120.8 209.64,-115.73\"/>\n",
       "</g>\n",
       "<!-- 8 -->\n",
       "<g id=\"node9\" class=\"node\">\n",
       "<title>8</title>\n",
       "<polygon fill=\"lightyellow\" stroke=\"none\" points=\"402.25,-105.75 275.75,-105.75 275.75,-73.25 402.25,-73.25 402.25,-105.75\"/>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"275.75,-73.25 275.75,-105.75 342.75,-105.75 342.75,-73.25 275.75,-73.25\"/>\n",
       "<text text-anchor=\"start\" x=\"280.75\" y=\"-91.25\" font-family=\"Linux libertine\" font-size=\"10.00\">output&#45;tensor</text>\n",
       "<text text-anchor=\"start\" x=\"293.12\" y=\"-80\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:0</text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"342.75,-73.25 342.75,-105.75 402.25,-105.75 402.25,-73.25 342.75,-73.25\"/>\n",
       "<text text-anchor=\"start\" x=\"347.75\" y=\"-85.62\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 10, 384)</text>\n",
       "</g>\n",
       "<!-- 5&#45;&gt;8 -->\n",
       "<g id=\"edge7\" class=\"edge\">\n",
       "<title>5&#45;&gt;8</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M279.28,-146.53C290.19,-136.3 303.52,-123.79 314.81,-113.2\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"316.98,-115.96 321.88,-106.57 312.19,-110.86 316.98,-115.96\"/>\n",
       "</g>\n",
       "<!-- 7 -->\n",
       "<g id=\"node8\" class=\"node\">\n",
       "<title>7</title>\n",
       "<polygon fill=\"lightyellow\" stroke=\"none\" points=\"233,-32.5 123,-32.5 123,0 233,0 233,-32.5\"/>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"123,0 123,-32.5 190,-32.5 190,0 123,0\"/>\n",
       "<text text-anchor=\"start\" x=\"128\" y=\"-18\" font-family=\"Linux libertine\" font-size=\"10.00\">output&#45;tensor</text>\n",
       "<text text-anchor=\"start\" x=\"140.38\" y=\"-6.75\" font-family=\"Linux libertine\" font-size=\"10.00\">depth:0</text>\n",
       "<polygon fill=\"none\" stroke=\"black\" points=\"190,0 190,-32.5 233,-32.5 233,0 190,0\"/>\n",
       "<text text-anchor=\"start\" x=\"195\" y=\"-12.38\" font-family=\"Linux libertine\" font-size=\"10.00\">(1, 384)</text>\n",
       "</g>\n",
       "<!-- 6&#45;&gt;7 -->\n",
       "<g id=\"edge6\" class=\"edge\">\n",
       "<title>6&#45;&gt;7</title>\n",
       "<path fill=\"none\" stroke=\"black\" d=\"M178,-68.65C178,-60.91 178,-51.91 178,-43.63\"/>\n",
       "<polygon fill=\"black\" stroke=\"black\" points=\"181.5,-43.76 178,-33.76 174.5,-43.76 181.5,-43.76\"/>\n",
       "</g>\n",
       "</g>\n",
       "</svg>\n"
      ],
      "text/plain": [
       "<graphviz.graphs.Digraph at 0x17c425dd0>"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from torch import LongTensor\n",
    "from torchview import draw_graph\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "\n",
    "model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')\n",
    "\n",
    "model_graph = draw_graph(model, input_size=(1, 10), dtypes=[LongTensor], depth=1)\n",
    "\n",
    "model_graph.visual_graph\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "d",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}