Updating from rollback

data_measurements/embeddings.py

@@ -20,12 +20,14 @@ import plotly.graph_objects as go
 import torch
 import transformers
 from datasets import load_from_disk
+from plotly.io import read_json
 from tqdm import tqdm
 
-from .dataset_utils import EMBEDDING_FIELD, OUR_TEXT_FIELD
+from .dataset_utils import EMBEDDING_FIELD
 
 
 def sentence_mean_pooling(model_output, attention_mask):
+    """Mean pooling of token embeddings for a sentence."""
     token_embeddings = model_output[
         0
     ]  # First element of model_output contains all token embeddings
@@ -38,46 +40,46 @@ def sentence_mean_pooling(model_output, attention_mask):
 
 
 class Embeddings:
-    def __init__(self, dstats, use_cache=False):
+    def __init__(
+        self,
+        dstats=None,
+        text_dset=None,
+        text_field_name="text",
+        cache_path="",
+        use_cache=False,
+    ):
         """Item embeddings and clustering"""
         self.device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        self.model_name = "sentence-transformers/all-mpnet-base-v2"
+        self.tokenizer = transformers.AutoTokenizer.from_pretrained(self.model_name)
+        self.model = transformers.AutoModel.from_pretrained(self.model_name).to(
+            self.device
+        )
+        self.text_dset = text_dset if dstats is None else dstats.text_dset
+        self.text_field_name = (
+            text_field_name if dstats is None else dstats.our_text_field
+        )
+        self.cache_path = cache_path if dstats is None else dstats.cache_path
+        self.embeddings_dset_fid = pjoin(self.cache_path, "embeddings_dset")
+        self.embeddings_dset = None
+        self.node_list_fid = pjoin(self.cache_path, "node_list.th")
         self.node_list = None
         self.nid_map = None
-        self.embeddings_dset = None
+        self.fig_tree_fid = pjoin(self.cache_path, "node_figure.json")
         self.fig_tree = None
         self.cached_clusters = {}
-        self.dstats = dstats
-        self.cache_path = dstats.cache_path
-        self.node_list_fid = pjoin(self.cache_path, "node_list.th")
         self.use_cache = use_cache
-        self.tokenizer = transformers.AutoTokenizer.from_pretrained(
-            "sentence-transformers/all-mpnet-base-v2"
-        )
-        self.model = transformers.AutoModel.from_pretrained(
-            "sentence-transformers/all-mpnet-base-v2"
-        ).to(self.device)
-
-    def make_text_embeddings(self):
-        embeddings_dset_fid = pjoin(self.cache_path, "embeddings_dset")
-        if self.use_cache and exists(embeddings_dset_fid):
-            self.embeddings_dset = load_from_disk(embeddings_dset_fid)
-        else:
-            self.embeddings_dset = self.make_embeddings()
-            self.embeddings_dset.save_to_disk(embeddings_dset_fid)
-
-    def make_hierarchical_clustering(self):
-        if self.use_cache and exists(self.node_list_fid):
-            self.node_list = torch.load(self.node_list_fid)
-        else:
-            self.make_text_embeddings()
-            self.node_list = self.fast_cluster(self.embeddings_dset, EMBEDDING_FIELD)
-            torch.save(self.node_list, self.node_list_fid)
-        self.nid_map = dict(
-            [(node["nid"], nid) for nid, node in enumerate(self.node_list)]
-        )
-        self.fig_tree = make_tree_plot(self.node_list, self.dstats.text_dset)
 
     def compute_sentence_embeddings(self, sentences):
+        """
+        Takes a list of sentences and computes their embeddings
+        using self.tokenizer and self.model (with output dimension D)
+        followed by mean pooling of the token representations and normalization
+        Args:
+            sentences ([string]): list of N input sentences
+        Returns:
+            torch.Tensor: sentence embeddings, dimension NxD
+        """
         batch = self.tokenizer(
             sentences, padding=True, truncation=True, return_tensors="pt"
         )
@@ -91,212 +93,70 @@ class Embeddings:
             return sentence_embeds
 
     def make_embeddings(self):
+        """
+        Batch computes the embeddings of the Dataset self.text_dset,
+        using the field self.text_field_name as input.
+        Returns:
+            Dataset: HF dataset object with a single EMBEDDING_FIELD field
+                corresponding to the embeddings (list of floats)
+        """
+
         def batch_embed_sentences(sentences):
             return {
                 EMBEDDING_FIELD: [
                     embed.tolist()
                     for embed in self.compute_sentence_embeddings(
-                        sentences[OUR_TEXT_FIELD]
+                        sentences[self.text_field_name]
                     )
                 ]
             }
 
-        text_dset_embeds = self.dstats.text_dset.map(
+        self.embeddings_dset = self.text_dset.map(
             batch_embed_sentences,
             batched=True,
             batch_size=32,
-            remove_columns=[self.dstats.our_text_field],
-        )
-
-        return text_dset_embeds
-
-    @staticmethod
-    def prepare_merges(embeddings, batch_size, low_thres=0.5):
-        top_idx_pre = torch.cat(
-            [torch.LongTensor(range(embeddings.shape[0]))[:, None]] * batch_size, dim=1
-        )
-        top_val_all = torch.Tensor(0, batch_size)
-        top_idx_all = torch.LongTensor(0, batch_size)
-        n_batches = math.ceil(len(embeddings) / batch_size)
-        for b in tqdm(range(n_batches)):
-            cos_scores = torch.mm(
-                embeddings[b * batch_size : (b + 1) * batch_size], embeddings.t()
-            )
-            for i in range(cos_scores.shape[0]):
-                cos_scores[i, (b * batch_size) + i :] = -1
-            top_val_large, top_idx_large = cos_scores.topk(
-                k=batch_size, dim=-1, largest=True
-            )
-            top_val_all = torch.cat([top_val_all, top_val_large], dim=0)
-            top_idx_all = torch.cat([top_idx_all, top_idx_large], dim=0)
-
-        all_merges = torch.cat(
-            [
-                top_idx_pre[top_val_all > low_thres][:, None],
-                top_idx_all[top_val_all > low_thres][:, None],
-            ],
-            dim=1,
+            remove_columns=[self.text_field_name],
         )
-        all_merge_scores = top_val_all[top_val_all > low_thres]
-        return (all_merges, all_merge_scores)
 
-    @staticmethod
-    def merge_nodes(nodes, current_thres, previous_thres, all_merges, all_merge_scores):
-        merge_ids = (all_merge_scores <= previous_thres) * (
-            all_merge_scores > current_thres
-        )
-        merges = all_merges[merge_ids]
-        for a, b in merges.tolist():
-            node_a = nodes[a]
-            while node_a["parent_id"] != -1:
-                node_a = nodes[node_a["parent_id"]]
-            node_b = nodes[b]
-            while node_b["parent_id"] != -1:
-                node_b = nodes[node_b["parent_id"]]
-            if node_a["nid"] == node_b["nid"]:
-                continue
-            else:
-                # merge if threshold allows
-                if (node_a["depth"] + node_b["depth"]) > 0 and min(
-                    node_a["merge_threshold"], node_b["merge_threshold"]
-                ) == current_thres:
-                    merge_to = None
-                    merge_from = None
-                    if node_a["nid"] < node_b["nid"]:
-                        merge_from = node_a
-                        merge_to = node_b
-                    if node_a["nid"] > node_b["nid"]:
-                        merge_from = node_b
-                        merge_to = node_a
-                    merge_to["depth"] = max(merge_to["depth"], merge_from["depth"])
-                    merge_to["weight"] += merge_from["weight"]
-                    merge_to["children_ids"] += (
-                        merge_from["children_ids"]
-                        if merge_from["depth"] > 0
-                        else [merge_from["nid"]]
-                    )
-                    for cid in merge_from["children_ids"]:
-                        nodes[cid]["parent_id"] = merge_to["nid"]
-                    merge_from["parent_id"] = merge_to["nid"]
-                # else new node
-                else:
-                    new_nid = len(nodes)
-                    new_node = {
-                        "nid": new_nid,
-                        "parent_id": -1,
-                        "depth": max(node_a["depth"], node_b["depth"]) + 1,
-                        "weight": node_a["weight"] + node_b["weight"],
-                        "children": [],
-                        "children_ids": [node_a["nid"], node_b["nid"]],
-                        "example_ids": [],
-                        "merge_threshold": current_thres,
-                    }
-                    node_a["parent_id"] = new_nid
-                    node_b["parent_id"] = new_nid
-                    nodes += [new_node]
-        return nodes
+        return self.embeddings_dset
 
-    def finalize_node(self, node, nodes, min_cluster_size):
-        node["children"] = sorted(
-            [
-                self.finalize_node(nodes[cid], nodes, min_cluster_size)
-                for cid in node["children_ids"]
-            ],
-            key=lambda x: x["weight"],
-            reverse=True,
-        )
-        if node["depth"] > 0:
-            node["example_ids"] = [
-                eid for child in node["children"] for eid in child["example_ids"]
-            ]
-        node["children"] = [
-            child for child in node["children"] if child["weight"] >= min_cluster_size
-        ]
-        assert node["weight"] == len(node["example_ids"]), print(node)
-        return node
+    def make_text_embeddings(self):
+        """Load embeddings dataset from cache or compute it."""
+        if self.use_cache and exists(self.embeddings_dset_fid):
+            self.embeddings_dset = load_from_disk(self.embeddings_dset_fid)
+        else:
+            self.embeddings_dset = self.make_embeddings()
+            self.embeddings_dset.save_to_disk(self.embeddings_dset_fid)
 
-    def fast_cluster(
+    def make_hierarchical_clustering(
         self,
-        text_dset_embeds,
-        embedding_field,
         batch_size=1000,
+        approx_neighbors=1000,
         min_cluster_size=10,
-        low_thres=0.5,
     ):
-        embeddings = torch.Tensor(text_dset_embeds[embedding_field])
-        batch_size = min(embeddings.shape[0], batch_size)
-        all_merges, all_merge_scores = self.prepare_merges(
-            embeddings, batch_size, low_thres
-        )
-        # prepare leaves
-        nodes = [
-            {
-                "nid": nid,
-                "parent_id": -1,
-                "depth": 0,
-                "weight": 1,
-                "children": [],
-                "children_ids": [],
-                "example_ids": [nid],
-                "merge_threshold": 1.0,
-            }
-            for nid in range(embeddings.shape[0])
-        ]
-        # one level per threshold range
-        for i in range(10):
-            p_thres = 1 - i * 0.05
-            c_thres = 0.95 - i * 0.05
-            nodes = self.merge_nodes(
-                nodes, c_thres, p_thres, all_merges, all_merge_scores
+        if self.use_cache and exists(self.node_list_fid):
+            self.node_list, self.nid_map = torch.load(self.node_list_fid)
+        else:
+            self.make_text_embeddings()
+            embeddings = torch.Tensor(self.embeddings_dset[EMBEDDING_FIELD])
+            self.node_list = fast_cluster(
+                embeddings, batch_size, approx_neighbors, min_cluster_size
             )
-        # make root
-        root_children = [
-            node
-            for node in nodes
-            if node["parent_id"] == -1 and node["weight"] >= min_cluster_size
-        ]
-        root = {
-            "nid": len(nodes),
-            "parent_id": -1,
-            "depth": max([node["depth"] for node in root_children]) + 1,
-            "weight": sum([node["weight"] for node in root_children]),
-            "children": [],
-            "children_ids": [node["nid"] for node in root_children],
-            "example_ids": [],
-            "merge_threshold": -1.0,
-        }
-        nodes += [root]
-        for node in root_children:
-            node["parent_id"] = root["nid"]
-        # finalize tree
-        tree = self.finalize_node(root, nodes, min_cluster_size)
-        node_list = []
-
-        def rec_map_nodes(node, node_list):
-            node_list += [node]
-            for child in node["children"]:
-                rec_map_nodes(child, node_list)
-
-        rec_map_nodes(tree, node_list)
-        # get centroids and distances
-        for node in node_list:
-            node_embeds = embeddings[node["example_ids"]]
-            node["centroid"] = node_embeds.sum(dim=0)
-            node["centroid"] /= node["centroid"].norm()
-            node["centroid_dot_prods"] = torch.mv(node_embeds, node["centroid"])
-            node["sorted_examples_centroid"] = sorted(
-                [
-                    (eid, edp.item())
-                    for eid, edp in zip(node["example_ids"], node["centroid_dot_prods"])
-                ],
-                key=lambda x: x[1],
-                reverse=True,
+            self.nid_map = dict(
+                [(node["nid"], nid) for nid, node in enumerate(self.node_list)]
+            )
+            torch.save((self.node_list, self.nid_map), self.node_list_fid)
+        if self.use_cache and exists(self.fig_tree_fid):
+            self.fig_tree = read_json(self.fig_tree_fid)
+        else:
+            self.fig_tree = make_tree_plot(
+                self.node_list, self.text_dset, self.text_field_name
             )
-        return node_list
+            self.fig_tree.write_json(self.fig_tree_fid)
 
     def find_cluster_beam(self, sentence, beam_size=20):
         """
-        This function finds the `beam_size` lef clusters that are closest to the
+        This function finds the `beam_size` leaf clusters that are closest to the
         proposed sentence and returns the full path from the root to the cluster
         along with the dot product between the sentence embedding and the
         cluster centroid
@@ -365,25 +225,268 @@ class Embeddings:
         )[:beam_size]
 
 
-def make_tree_plot(node_list, text_dset):
+def prepare_merges(embeddings, batch_size=1000, approx_neighbors=1000, low_thres=0.5):
+    """
+    Prepares an initial list of merges for hierarchical
+    clustering. First compute the `approx_neighbors` nearest neighbors,
+    then propose a merge for any two points that are closer than `low_thres`
+
+    Note that if a point has more than `approx_neighbors` neighbors
+    closer than `low_thres`, this approach will miss some of those merges
+
+    Args:
+        embeddings (toch.Tensor): Tensor of sentence embeddings - dimension NxD
+        batch_size (int): compute nearest neighbors of `batch_size` points at a time
+        approx_neighbors (int): only keep `approx_neighbors` nearest neighbors of a point
+        low_thres (float): only return merges where the dot product is greater than `low_thres`
+    Returns:
+        torch.LongTensor: proposed merges ([i, j] with i>j) - dimension: Mx2
+        torch.Tensor: merge scores - dimension M
+    """
+    top_idx_pre = torch.cat(
+        [torch.LongTensor(range(embeddings.shape[0]))[:, None]] * batch_size, dim=1
+    )
+    top_val_all = torch.Tensor(0, approx_neighbors)
+    top_idx_all = torch.LongTensor(0, approx_neighbors)
+    n_batches = math.ceil(len(embeddings) / batch_size)
+    for b in tqdm(range(n_batches)):
+        # TODO: batch across second dimension
+        cos_scores = torch.mm(
+            embeddings[b * batch_size : (b + 1) * batch_size], embeddings.t()
+        )
+        for i in range(cos_scores.shape[0]):
+            cos_scores[i, (b * batch_size) + i :] = -1
+        top_val_large, top_idx_large = cos_scores.topk(
+            k=approx_neighbors, dim=-1, largest=True
+        )
+        top_val_all = torch.cat([top_val_all, top_val_large], dim=0)
+        top_idx_all = torch.cat([top_idx_all, top_idx_large], dim=0)
+        max_neighbor_dist = top_val_large[:, -1].max().item()
+        if max_neighbor_dist > low_thres:
+            print(
+                f"WARNING: with the current set of neireast neighbor, the farthest is {max_neighbor_dist}"
+            )
+
+    all_merges = torch.cat(
+        [
+            top_idx_pre[top_val_all > low_thres][:, None],
+            top_idx_all[top_val_all > low_thres][:, None],
+        ],
+        dim=1,
+    )
+    all_merge_scores = top_val_all[top_val_all > low_thres]
+
+    return (all_merges, all_merge_scores)
+
+
+def merge_nodes(nodes, current_thres, previous_thres, all_merges, all_merge_scores):
+    """
+    Merge all nodes if the max dot product between any of their descendants
+    is greater than current_thres.
+
+    Args:
+        nodes ([dict]): list of dicts representing the current set of nodes
+        current_thres (float): merge all nodes closer than current_thres
+        previous_thres (float): nodes closer than previous_thres are already merged
+        all_merges (torch.LongTensor): proposed merges ([i, j] with i>j) - dimension: Mx2
+        all_merge_scores (torch.Tensor): merge scores - dimension M
+    Returns:
+        [dict]: extended list with the newly created internal nodes
+    """
+    merge_ids = (all_merge_scores <= previous_thres) * (
+        all_merge_scores > current_thres
+    )
+    if merge_ids.sum().item() > 0:
+        merges = all_merges[merge_ids]
+        for a, b in merges.tolist():
+            node_a = nodes[a]
+            while node_a["parent_id"] != -1:
+                node_a = nodes[node_a["parent_id"]]
+            node_b = nodes[b]
+            while node_b["parent_id"] != -1:
+                node_b = nodes[node_b["parent_id"]]
+            if node_a["nid"] == node_b["nid"]:
+                continue
+            else:
+                # merge if threshold allows
+                if (node_a["depth"] + node_b["depth"]) > 0 and min(
+                    node_a["merge_threshold"], node_b["merge_threshold"]
+                ) == current_thres:
+                    merge_to = None
+                    merge_from = None
+                    if node_a["nid"] < node_b["nid"]:
+                        merge_from = node_a
+                        merge_to = node_b
+                    if node_a["nid"] > node_b["nid"]:
+                        merge_from = node_b
+                        merge_to = node_a
+                    merge_to["depth"] = max(merge_to["depth"], merge_from["depth"])
+                    merge_to["weight"] += merge_from["weight"]
+                    merge_to["children_ids"] += (
+                        merge_from["children_ids"]
+                        if merge_from["depth"] > 0
+                        else [merge_from["nid"]]
+                    )
+                    for cid in merge_from["children_ids"]:
+                        nodes[cid]["parent_id"] = merge_to["nid"]
+                    merge_from["parent_id"] = merge_to["nid"]
+                # else new node
+                else:
+                    new_nid = len(nodes)
+                    new_node = {
+                        "nid": new_nid,
+                        "parent_id": -1,
+                        "depth": max(node_a["depth"], node_b["depth"]) + 1,
+                        "weight": node_a["weight"] + node_b["weight"],
+                        "children": [],
+                        "children_ids": [node_a["nid"], node_b["nid"]],
+                        "example_ids": [],
+                        "merge_threshold": current_thres,
+                    }
+                    node_a["parent_id"] = new_nid
+                    node_b["parent_id"] = new_nid
+                    nodes += [new_node]
+    return nodes
+
+
+def finalize_node(node, nodes, min_cluster_size):
+    """Post-process nodes to sort children by descending weight,
+    get full list of leaves in the sub-tree, and direct links
+    to the cildren nodes, then recurses to all children.
+
+    Nodes with fewer than `min_cluster_size` descendants are collapsed
+    into a single leaf.
+    """
+    node["children"] = sorted(
+        [
+            finalize_node(nodes[cid], nodes, min_cluster_size)
+            for cid in node["children_ids"]
+        ],
+        key=lambda x: x["weight"],
+        reverse=True,
+    )
+    if node["depth"] > 0:
+        node["example_ids"] = [
+            eid for child in node["children"] for eid in child["example_ids"]
+        ]
+    node["children"] = [
+        child for child in node["children"] if child["weight"] >= min_cluster_size
+    ]
+    assert node["weight"] == len(node["example_ids"]), print(node)
+    return node
+
+
+def fast_cluster(
+    embeddings,
+    batch_size=1000,
+    approx_neighbors=1000,
+    min_cluster_size=10,
+    low_thres=0.5,
+):
+    """
+    Computes an approximate hierarchical clustering based on example
+    embeddings. The join criterion is min clustering, i.e. two clusters
+    are joined if any pair of their descendants are closer than a threshold
+
+    The approximate comes from the fact that only the `approx_neighbors` nearest
+    neighbors of an example are considered for merges
+    """
+    batch_size = min(embeddings.shape[0], batch_size)
+    all_merges, all_merge_scores = prepare_merges(
+        embeddings, batch_size, approx_neighbors, low_thres
+    )
+    # prepare leaves
+    nodes = [
+        {
+            "nid": nid,
+            "parent_id": -1,
+            "depth": 0,
+            "weight": 1,
+            "children": [],
+            "children_ids": [],
+            "example_ids": [nid],
+            "merge_threshold": 1.0,
+        }
+        for nid in range(embeddings.shape[0])
+    ]
+    # one level per threshold range
+    for i in range(10):
+        p_thres = 1 - i * 0.05
+        c_thres = 0.95 - i * 0.05
+        nodes = merge_nodes(nodes, c_thres, p_thres, all_merges, all_merge_scores)
+    # make root
+    root_children = [
+        node
+        for node in nodes
+        if node["parent_id"] == -1 and node["weight"] >= min_cluster_size
+    ]
+    root = {
+        "nid": len(nodes),
+        "parent_id": -1,
+        "depth": max([node["depth"] for node in root_children]) + 1,
+        "weight": sum([node["weight"] for node in root_children]),
+        "children": [],
+        "children_ids": [node["nid"] for node in root_children],
+        "example_ids": [],
+        "merge_threshold": -1.0,
+    }
+    nodes += [root]
+    for node in root_children:
+        node["parent_id"] = root["nid"]
+    # finalize tree
+    tree = finalize_node(root, nodes, min_cluster_size)
+    node_list = []
+
+    def rec_map_nodes(node, node_list):
+        node_list += [node]
+        for child in node["children"]:
+            rec_map_nodes(child, node_list)
+
+    rec_map_nodes(tree, node_list)
+    # get centroids and distances
+    for node in node_list:
+        node_embeds = embeddings[node["example_ids"]]
+        node["centroid"] = node_embeds.sum(dim=0)
+        node["centroid"] /= node["centroid"].norm()
+        node["centroid_dot_prods"] = torch.mv(node_embeds, node["centroid"])
+        node["sorted_examples_centroid"] = sorted(
+            [
+                (eid, edp.item())
+                for eid, edp in zip(node["example_ids"], node["centroid_dot_prods"])
+            ],
+            key=lambda x: x[1],
+            reverse=True,
+        )
+    return node_list
+
+
+def make_tree_plot(node_list, text_dset, text_field_name):
+    """
+    Makes a graphical representation of the tree encoded
+    in node-list. The hover label for each node shows the number
+    of descendants and the 5 examples that are closest to the centroid
+    """
     nid_map = dict([(node["nid"], nid) for nid, node in enumerate(node_list)])
 
     for nid, node in enumerate(node_list):
+        # get list of
+        node_examples = {}
+        for sid, score in node["sorted_examples_centroid"]:
+            node_examples[text_dset[sid][text_field_name]] = score
+            if len(node_examples) >= 5:
+                break
         node["label"] = node.get(
             "label",
             f"{nid:2d} - {node['weight']:5d} items <br>"
             + "<br>".join(
                 [
-                    "> " + txt[:64] + ("..." if len(txt) >= 63 else "")
-                    for txt in list(
-                        set(text_dset.select(node["example_ids"])[OUR_TEXT_FIELD])
-                    )[:5]
+                    f" {score:.2f} > {txt[:64]}" + ("..." if len(txt) >= 63 else "")
+                    for txt, score in node_examples.items()
                 ]
             ),
         )
 
     # make plot nodes
-    # TODO: something more efficient than set to remove duplicates
     labels = [node["label"] for node in node_list]
 
     root = node_list[0]

data_measurements/streamlit_utils.py

@@ -21,7 +21,6 @@ from st_aggrid import AgGrid, GridOptionsBuilder
 
 from .dataset_utils import HF_DESC_FIELD, HF_FEATURE_FIELD, HF_LABEL_FIELD
 
-
 def sidebar_header():
     st.sidebar.markdown(
         """
@@ -167,7 +166,11 @@ def expander_text_lengths(dstats, column_id):
         st.markdown(
             "### Here is the relative frequency of different text lengths in your dataset:"
         )
-        st.plotly_chart(dstats.fig_tok_length, use_container_width=True)
+        #TODO: figure out more elegant way to do this:
+        try:
+            st.image(dstats.fig_tok_length_png)
+        except:
+            st.pyplot(dstats.fig_tok_length, use_container_width=True)
         st.markdown(
             "The average length of text instances is **"
             + str(dstats.avg_length)
@@ -175,19 +178,11 @@ def expander_text_lengths(dstats, column_id):
             + str(dstats.std_length)
             + "**."
         )
-
-        start_id_show_lengths = st.slider(
-            f"Show the shortest sentences{column_id} starting at:",
-            0,
-            dstats.num_uniq_lengths,
-            value=0,
-            step=1,
-        )
-
         # This is quite a large file and is breaking our ability to navigate the app development.
         # Just passing if it's not already there for launch v0
         if dstats.length_df is not None:
-            st.dataframe(dstats.length_df[dstats.length_df["length"] == start_id_show_lengths].set_index("length"))
+            start_id_show_lengths= st.selectbox("Show examples of length:", sorted(dstats.length_df["length"].unique().tolist()))
+            st.table(dstats.length_df[dstats.length_df["length"] == start_id_show_lengths].set_index("length"))
 
 
 ### Third, use a sentence embedding model
@@ -285,17 +280,7 @@ def expander_text_duplicates(dstats, column_id):
         if dstats.dup_counts_df is None or dstats.dup_counts_df.empty:
             st.write("There are no duplicates in this dataset! 🥳")
         else:
-            gb = GridOptionsBuilder.from_dataframe(dstats.dup_counts_df)
-            gb.configure_column(
-                f"text{column_id}",
-                wrapText=True,
-                resizable=True,
-                autoHeight=True,
-                min_column_width=85,
-                use_container_width=True,
-            )
-            go = gb.build()
-            AgGrid(dstats.dup_counts_df, gridOptions=go)
+            st.dataframe(dstats.dup_counts_df.reset_index(drop=True))
 
 
 def expander_npmi_description(min_vocab):