Update hierarchy inference notebook

penpot · Sep 25, 2024 · f8404b2 · f8404b2
1 parent 716b44d
commit f8404b2
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 14 deletions.
diff --git a/notebooks/hierarchy_inference.ipynb b/notebooks/hierarchy_inference.ipynb
@@ -18,7 +18,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 7,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-07-10T08:19:19.147150500Z",
@@ -34,9 +34,13 @@
     "    InteractiveHTMLHierarchyVisualizer,\n",
     "    InteractiveSVGHierarchyVisualizer,\n",
     ")\n",
+    "from penai.llm.llm_model import RegisteredLLM\n",
     "from penai.registries.projects import SavedPenpotProject\n",
     "from penai.utils.ipython import IFrameFromSrc\n",
-    "from penai.utils.vis import DesignElementVisualizer, ShapeHierarchyVisualizer"
+    "from penai.utils.vis import (\n",
+    "    DesignElementVisualizer,\n",
+    "    ShapeHighlighter,\n",
+    ")"
    ]
   },
   {
@@ -54,7 +58,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": 3,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-07-10T08:19:30.149660900Z",
@@ -66,7 +70,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Scanning remote paths in penpot/data/raw/designs/Material Design 3: 100%|██████████| 36/36 [00:00<00:00, 322.66it/s]\n",
+      "Scanning remote paths in penpot/data/raw/designs/Material Design 3: 100%|██████████| 36/36 [00:00<00:00, 328.15it/s]\n",
       "force pulling (bytes): 0it [00:00, ?it/s]\n"
      ]
     }
@@ -85,7 +89,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": 4,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-07-10T08:19:40.815981400Z",
@@ -97,7 +101,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "Setting view boxes: 100%|██████████| 163/163 [00:01<00:00, 100.61it/s]\n"
+      "Setting view boxes: 100%|██████████| 163/163 [00:02<00:00, 79.34it/s] \n"
      ]
     }
    ],
@@ -115,7 +119,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 7,
+   "execution_count": 5,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-07-10T08:19:40.974091600Z",
@@ -136,7 +140,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 6,
    "metadata": {
     "ExecuteTime": {
      "end_time": "2024-07-10T08:20:42.204256700Z",
@@ -145,11 +149,12 @@
    },
    "outputs": [],
    "source": [
-    "shape_visualizer = ShapeHierarchyVisualizer()\n",
+    "shape_visualizer = ShapeHighlighter()\n",
     "design_element_visualizer = DesignElementVisualizer(shape_visualizer=shape_visualizer)\n",
     "\n",
     "hierarchy_inference = HierarchyInferencer(\n",
     "    shape_visualizer=design_element_visualizer,\n",
+    "    model=RegisteredLLM.CLAUDE_3_5_SONNET\n",
     ")"
    ]
   },
@@ -164,9 +169,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 163/163 [00:00<00:00, 4299821.08it/s]\n",
+      "Scanning remote paths in penpot/data/cache/llm_responses_cache.sqlite: 100%|██████████| 1/1 [00:00<00:00, 16.60it/s]\n",
+      "force pulling (bytes): 100%|██████████| 2465792/2465792 [00:00<00:00, 10025456.76it/s]\n"
+     ]
+    }
+   ],
    "source": [
     "hierarchy = hierarchy_inference.infer_shape_hierarchy(cover_frame)"
    ]

diff --git a/src/penai/hierarchy_generation/utils.py b/src/penai/hierarchy_generation/utils.py
@@ -273,9 +273,6 @@ def __init__(
         svg_path: str | None = None,
         title: str = "Hierarchy Inspection",
     ):
-        with open(os.path.join(top_level_directory, "resources", "hierarchy.html")) as f:
-            html_content = f.read()
-
         assert count_not_none(svg, svg_path) == 1
 
         jstree_data_dict = self._create_jstree_data_dict(hierarchy_element)