interactive_simplicial_complex_from_texts.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "116a1ca2-cd11-4cd6-b1a9-ecfcf97f139e",
   "metadata": {},
   "source": [
    "# Interactive Simplicial Complex from Text"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae60a9d4-a2d7-41f9-b475-e15dfc4288e8",
   "metadata": {},
   "source": [
    "An interactive simplicial complex from persistent homology of an information theoretic analysis of attention. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "a5d5a870-1a8b-434d-9a6a-8761f2b3db57",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "40a59886276f464d825d24f379aedccd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.05, continuous_update=False, description='Threshold:', max=0.4, min=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.update_plot(threshold)>"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from transformers import GPT2Tokenizer, GPT2Model\n",
    "import gudhi as gd\n",
    "import networkx as nx\n",
    "from scipy.spatial.distance import jensenshannon\n",
    "import plotly.graph_objs as go\n",
    "import ipywidgets as widgets\n",
    "\n",
    "\n",
    "def get_attention_matrix(text, model, tokenizer, layer, head):\n",
    "    inputs = tokenizer(text, return_tensors=\"pt\")\n",
    "    outputs = model(**inputs, output_attentions=True)\n",
    "    attention = outputs.attentions[layer][0, head].detach().cpu().numpy()\n",
    "    return attention\n",
    "\n",
    "\n",
    "def compute_persistence(attention_matrix):\n",
    "    softmax_attention = np.exp(attention_matrix) / np.sum(np.exp(attention_matrix), axis=-1)[:, np.newaxis]\n",
    "    distance_matrix = np.array([[np.sqrt(jensenshannon(softmax_attention[i], softmax_attention[j])) for j in range(softmax_attention.shape[0])] for i in range(softmax_attention.shape[0])])\n",
    "\n",
    "    rips_complex = gd.RipsComplex(distance_matrix=distance_matrix, max_edge_length=np.inf)\n",
    "    simplex_tree = rips_complex.create_simplex_tree(max_dimension=2)\n",
    "    persistence = simplex_tree.persistence(min_persistence=0.01)\n",
    "    return persistence, simplex_tree, distance_matrix\n",
    "\n",
    "\n",
    "def plot_simplicial_complex_3d(simplex_tree, distance_matrix, title, threshold, tokens):\n",
    "    g = nx.Graph()\n",
    "    for (simplex, _) in simplex_tree.get_filtration():\n",
    "        if len(simplex) == 2:\n",
    "            if distance_matrix[simplex[0]][simplex[1]] <= threshold:\n",
    "                g.add_edge(simplex[0], simplex[1])\n",
    "\n",
    "    labels = {node: tokens[node] for node in g.nodes()}\n",
    "    pos = nx.spring_layout(g, dim=3, seed=42)\n",
    "    Xn = [pos[k][0] for k in g.nodes()]\n",
    "    Yn = [pos[k][1] for k in g.nodes()]\n",
    "    Zn = [pos[k][2] for k in g.nodes()]\n",
    "\n",
    "    Xe = []\n",
    "    Ye = []\n",
    "    Ze = []\n",
    "    for e in g.edges():\n",
    "        Xe += [pos[e[0]][0], pos[e[1]][0], None]\n",
    "        Ye += [pos[e[0]][1], pos[e[1]][1], None]\n",
    "        Ze += [pos[e[0]][2], pos[e[1]][2], None]\n",
    "\n",
    "    trace_edges = go.Scatter3d(x=Xe, y=Ye, z=Ze, mode='lines', line=dict(color='gray', width=1))\n",
    "    trace_nodes = go.Scatter3d(x=Xn, y=Yn, z=Zn, mode='markers+text', text=list(labels.values()), marker=dict(symbol='circle', size=10, color='lightblue'), textposition=\"top center\")\n",
    "\n",
    "    layout = go.Layout(title=title, scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'), showlegend=False)\n",
    "    fig = go.Figure(data=[trace_edges, trace_nodes], layout=layout)\n",
    "    fig.show()\n",
    "\n",
    "\n",
    "if __name__ == '__main__':\n",
    "    tokenizer = GPT2Tokenizer.from_pretrained('gpt2')\n",
    "    model = GPT2Model.from_pretrained('gpt2')\n",
    "\n",
    "    \n",
    "text1 = \"Quantum information theory is interesting\"\n",
    "text2 = \"Quantum information theory allows us to study attention using entanglement\"\n",
    "\n",
    "layer = 1\n",
    "head = 2\n",
    "\n",
    "attention_matrix1 = get_attention_matrix(text1, model, tokenizer, layer, head)\n",
    "attention_matrix2 = get_attention_matrix(text2, model, tokenizer, layer, head)\n",
    "\n",
    "persistence1, simplex_tree1, distance_matrix1 = compute_persistence(attention_matrix1)\n",
    "persistence2, simplex_tree2, distance_matrix2 = compute_persistence(attention_matrix2)\n",
    "\n",
    "tokens1 = [tokenizer.decode(token_id) for token_id in tokenizer.encode(text1)]\n",
    "tokens2 = [tokenizer.decode(token_id) for token_id in tokenizer.encode(text2)]\n",
    "\n",
    "threshold_slider = widgets.FloatSlider(\n",
    "    value=0.05,\n",
    "    min=0.01,\n",
    "    max=0.4,\n",
    "    step=0.01,\n",
    "    description='Threshold:',\n",
    "    continuous_update=False\n",
    ")\n",
    "\n",
    "def update_plot(threshold):\n",
    "    plot_simplicial_complex_3d(simplex_tree1, distance_matrix1, \"Simplicial Complex for Text 1\", threshold, tokens1)\n",
    "    plot_simplicial_complex_3d(simplex_tree2, distance_matrix2, \"Simplicial Complex for Text 2\", threshold, tokens2)\n",
    "\n",
    "widgets.interact(update_plot, threshold=threshold_slider)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2f31d3a7-2d50-46b4-b91b-4c9fdb6c307f",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-multilingual-cased were not used when initializing BertModel: ['cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.bias', 'cls.seq_relationship.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.bias']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "56ecd2400179431e8e546ed25cf1ba13",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.37, continuous_update=False, description='Threshold:', max=0.6, min=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.update_plot(threshold)>"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from transformers import GPT2Tokenizer, GPT2Model\n",
    "import gudhi as gd\n",
    "import matplotlib.pyplot as plt\n",
    "import networkx as nx\n",
    "from scipy.spatial.distance import jensenshannon\n",
    "import plotly.graph_objs as go\n",
    "import ipywidgets as widgets\n",
    "from IPython.display import display\n",
    "\n",
    "\n",
    "def get_attention_matrix(text, model, tokenizer, layer, head):\n",
    "    inputs = tokenizer(text, return_tensors=\"pt\")\n",
    "    outputs = model(**inputs, output_attentions=True)\n",
    "    attention = outputs.attentions[layer][0, head].detach().cpu().numpy()\n",
    "    return attention\n",
    "\n",
    "def compute_persistence(attention_matrix):\n",
    "    softmax_attention = np.exp(attention_matrix) / np.sum(np.exp(attention_matrix), axis=-1)[:, np.newaxis]\n",
    "    distance_matrix = np.array([[np.sqrt(jensenshannon(softmax_attention[i], softmax_attention[j])) for j in range(softmax_attention.shape[0])] for i in range(softmax_attention.shape[0])])\n",
    "    \n",
    "    rips_complex = gd.RipsComplex(distance_matrix=distance_matrix, max_edge_length=np.inf)\n",
    "    simplex_tree = rips_complex.create_simplex_tree(max_dimension=2)\n",
    "    persistence = simplex_tree.persistence(min_persistence=0.01)\n",
    "    return persistence, simplex_tree, distance_matrix\n",
    "\n",
    "def plot_simplicial_complex_3d(simplex_tree, distance_matrix, title, threshold, tokens):\n",
    "    g = nx.Graph()\n",
    "    for (simplex, _) in simplex_tree.get_filtration():\n",
    "        if len(simplex) == 2:\n",
    "            if distance_matrix[simplex[0]][simplex[1]] <= threshold:\n",
    "                g.add_edge(simplex[0], simplex[1])\n",
    "\n",
    "    # Create a token dictionary only for nodes in the graph\n",
    "    labels = {node: tokens[node] for node in g.nodes()}\n",
    "    \n",
    "    # 3D layout\n",
    "    pos = nx.spring_layout(g, dim=3, seed=42)\n",
    "    \n",
    "    # Extract node coordinates\n",
    "    Xn = [pos[k][0] for k in g.nodes()]\n",
    "    Yn = [pos[k][1] for k in g.nodes()]\n",
    "    Zn = [pos[k][2] for k in g.nodes()]\n",
    "    \n",
    "    # Extract edge coordinates\n",
    "    Xe = []\n",
    "    Ye = []\n",
    "    Ze = []\n",
    "    for e in g.edges():\n",
    "        Xe += [pos[e[0]][0], pos[e[1]][0], None]\n",
    "        Ye += [pos[e[0]][1], pos[e[1]][1], None]\n",
    "        Ze += [pos[e[0]][2], pos[e[1]][2], None]\n",
    "    \n",
    "    # Create a trace for edges\n",
    "    trace_edges = go.Scatter3d(x=Xe, y=Ye, z=Ze, mode='lines', line=dict(color='gray', width=1))\n",
    "    \n",
    "    # Create a trace for nodes\n",
    "    trace_nodes = go.Scatter3d(x=Xn, y=Yn, z=Zn, mode='markers+text', text=list(labels.values()), marker=dict(symbol='circle', size=10, color='lightblue'), textposition=\"top center\")\n",
    "    \n",
    "    # Create a layout\n",
    "    layout = go.Layout(title=title, scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'), showlegend=False)\n",
    "    \n",
    "    # Create a plot\n",
    "    fig = go.Figure(data=[trace_edges, trace_nodes], layout=layout)\n",
    "    fig.show()\n",
    "\n",
    "\n",
    "from transformers import BertTokenizer, BertModel\n",
    "\n",
    "# Load pre-trained transformer model and tokenizer\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-multilingual-cased')\n",
    "model = BertModel.from_pretrained('bert-base-multilingual-cased')\n",
    "\n",
    "# Text inputs\n",
    "text1 = \"Quantum information theory is interesting\"\n",
    "text2 = \"תיאוריה קוונטית של מידע מעניינת\"  # Hebrew translation: \"Quantum information theory is interesting\"\n",
    "\n",
    "\n",
    "# Choose a layer and head\n",
    "layer = 1\n",
    "head = 2\n",
    "\n",
    "# Get attention matrices\n",
    "attention_matrix1 = get_attention_matrix(text1, model, tokenizer, layer, head)\n",
    "attention_matrix2 = get_attention_matrix(text2, model, tokenizer, layer, head)\n",
    "\n",
    "# Compute persistence and simplex trees\n",
    "persistence1, simplex_tree1, distance_matrix1 = compute_persistence(attention_matrix1)\n",
    "persistence2, simplex_tree2, distance_matrix2 = compute_persistence(attention_matrix2)\n",
    "\n",
    "\n",
    "# Get tokens\n",
    "tokens1 = tokenizer.tokenize(text1)\n",
    "tokens2 = tokenizer.tokenize(text2)\n",
    "\n",
    "# Add special tokens ([CLS], [SEP], etc.) if necessary\n",
    "tokens1 = [tokenizer.decode(token_id) for token_id in tokenizer.encode(text1)]\n",
    "tokens2 = [tokenizer.decode(token_id) for token_id in tokenizer.encode(text2)]\n",
    "\n",
    "# Create a slider for the threshold value\n",
    "threshold_slider = widgets.FloatSlider(\n",
    "    value=0.37,\n",
    "    min=0.37,\n",
    "    max=0.6,\n",
    "    step=0.001,\n",
    "    description='Threshold:',\n",
    "    continuous_update=False\n",
    ")\n",
    "\n",
    "# Define a function to update the plot based on the threshold value from the slider\n",
    "def update_plot(threshold):\n",
    "    plot_simplicial_complex_3d(simplex_tree1, distance_matrix1, \"Simplicial Complex for Text 1\", threshold, tokens1)\n",
    "    plot_simplicial_complex_3d(simplex_tree2, distance_matrix2, \"Simplicial Complex for Text 2\", threshold, tokens2)\n",
    "\n",
    "# Connect the slider to the update_plot function\n",
    "widgets.interact(update_plot, threshold=threshold_slider)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50e20063-95f8-4d45-bcb0-6baf3a7853de",
   "metadata": {},
   "source": [
    "Here a slightly different version that handles plots where there are no edges connecting nodes (that is, the $1$-skeleton of the complicial "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0ba833c9-6f90-4482-b923-e7fd87ad9f1a",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at bert-base-multilingual-cased were not used when initializing BertModel: ['cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.bias', 'cls.seq_relationship.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.bias']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d47d04474a5d44bba4c378a43d036626",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "interactive(children=(FloatSlider(value=0.37, continuous_update=False, description='Threshold:', max=0.6, min=…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "<function __main__.update_plot(threshold)>"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from transformers import BertTokenizer, BertModel\n",
    "import gudhi as gd\n",
    "import networkx as nx\n",
    "from scipy.spatial.distance import jensenshannon\n",
    "import plotly.graph_objs as go\n",
    "import ipywidgets as widgets\n",
    "\n",
    "\n",
    "def get_attention_matrix(text, model, tokenizer, layer, head):\n",
    "    inputs = tokenizer(text, return_tensors=\"pt\")\n",
    "    outputs = model(**inputs, output_attentions=True)\n",
    "    attention = outputs.attentions[layer][0, head].detach().cpu().numpy()\n",
    "    return attention\n",
    "\n",
    "\n",
    "def compute_persistence(attention_matrix):\n",
    "    softmax_attention = np.exp(attention_matrix) / np.sum(np.exp(attention_matrix), axis=-1, keepdims=True)\n",
    "    distance_matrix = np.array([[np.sqrt(jensenshannon(softmax_attention[i], softmax_attention[j])) for j in range(softmax_attention.shape[0])] for i in range(softmax_attention.shape[0])])\n",
    "\n",
    "    rips_complex = gd.RipsComplex(distance_matrix=distance_matrix, max_edge_length=np.inf)\n",
    "    simplex_tree = rips_complex.create_simplex_tree(max_dimension=2)\n",
    "    persistence = simplex_tree.persistence(min_persistence=0.01)\n",
    "    return persistence, simplex_tree, distance_matrix\n",
    "\n",
    "\n",
    "def plot_simplicial_complex_3d(simplex_tree, distance_matrix, title, threshold, tokens):\n",
    "    g = nx.Graph()\n",
    "    for (simplex, _) in simplex_tree.get_filtration():\n",
    "        if len(simplex) == 2:\n",
    "            if distance_matrix[simplex[0]][simplex[1]] <= threshold:\n",
    "                g.add_edge(simplex[0], simplex[1])\n",
    "\n",
    "    if not g.edges():\n",
    "        print(f\"No edges in the graph for '{title}'. Try adjusting the threshold value.\")\n",
    "        return\n",
    "\n",
    "    labels = {node: tokens[node] for node in g.nodes()}\n",
    "    pos = nx.spring_layout(g, dim=3, seed=42)\n",
    "    Xn, Yn, Zn = zip(*pos.values())\n",
    "    Xe, Ye, Ze = [], [], []\n",
    "    for e in g.edges():\n",
    "        Xe += [pos[e[0]][0], pos[e[1]][0], None]\n",
    "        Ye += [pos[e[0]][1], pos[e[1]][1], None]\n",
    "        Ze += [pos[e[0]][2], pos[e[1]][2], None]\n",
    "\n",
    "    trace_edges = go.Scatter3d(x=Xe, y=Ye, z=Ze, mode='lines', line=dict(color='gray', width=1))\n",
    "    trace_nodes = go.Scatter3d(x=Xn, y=Yn, z=Zn, mode='markers+text', text=list(labels.values()), marker=dict(symbol='circle', size=10, color='lightblue'), textposition=\"top center\")\n",
    "    layout = go.Layout(title=title, scene=dict(xaxis_title='X', yaxis_title='Y', zaxis_title='Z'), showlegend=False)\n",
    "    fig = go.Figure(data=[trace_edges, trace_nodes], layout=layout)\n",
    "    fig.show()\n",
    "\n",
    "\n",
    "tokenizer = BertTokenizer.from_pretrained('bert-base-multilingual-cased')\n",
    "model = BertModel.from_pretrained('bert-base-multilingual-cased')\n",
    "\n",
    "text1 = \"Quantum information theory is interesting\"\n",
    "text2 = \"תיאוריה קוונטית של מידע מעניינת\"\n",
    "\n",
    "layer = 1\n",
    "head = 2\n",
    "\n",
    "attention_matrix1 = get_attention_matrix(text1, model, tokenizer, layer, head)\n",
    "attention_matrix2 = get_attention_matrix(text2, model, tokenizer, layer, head)\n",
    "\n",
    "persistence1, simplex_tree1, distance_matrix1 = compute_persistence(attention_matrix1)\n",
    "persistence2, simplex_tree2, distance_matrix2 = compute_persistence(attention_matrix2)\n",
    "\n",
    "tokens1 = [tokenizer.decode(token_id) for token_id in tokenizer.encode(text1)]\n",
    "tokens2 = [tokenizer.decode(token_id) for token_id in tokenizer.encode(text2)]\n",
    "\n",
    "threshold_slider = widgets.FloatSlider(\n",
    "value=0.37,\n",
    "min=0.37,\n",
    "max=0.6,\n",
    "step=0.001,\n",
    "description='Threshold:',\n",
    "continuous_update=False\n",
    ")\n",
    "\n",
    "def update_plot(threshold):\n",
    "    plot_simplicial_complex_3d(simplex_tree1, distance_matrix1, \"Simplicial Complex for Text 1\", threshold, tokens1)\n",
    "    plot_simplicial_complex_3d(simplex_tree2, distance_matrix2, \"Simplicial Complex for Text 2\", threshold, tokens2)\n",
    "\n",
    "widgets.interact(update_plot, threshold=threshold_slider)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98ff2229-cf1c-4b92-8300-afe2f5e28d76",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.8.5 64-bit ('anaconda3': virtualenv)",
   "language": "python",
   "name": "python385jvsc74a57bd0474c67ce7e36ad5731492349411c4ce02ca5c170a680b2d1efe1eb0325e35fe7"
  },
  "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}