Repository: dirty-cat/dirty_cat (python)

Issue:
Help users find relevant ken types
We should provide a helper function to help people tune the "type" argument, maybe a function called "get_ken_matching_types" that would return all the types that have the corresponding substring. This function should be exposed to users and mentioned in the relevant examples and related functions.

Complement to #487

Pull request description:
Add functions for exploring KEN embeddings
Fixes #495

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index 1f13963..15a03e4 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -8,7 +8,9 @@ Release 0.4.1
 Major changes
 -------------
 * :func:`fuzzy_join` and :class:`FeatureAugmenter` now perform joins on missing values as in `pandas.merge`
-but raises a warning. :pr:`522` and :pr:`529` by :user:`Jovan Stojanovic <jovan-stojanovic>`
+  but raises a warning. :pr:`522` and :pr:`529` by :user:`Jovan Stojanovic <jovan-stojanovic>`
+* Added :func:`get_ken_table_aliases` and :func:`get_ken_types` for exploring
+  KEN embeddings. :pr:`539` by :user:`Lilian Boulard`.
 
 Minor changes
 -------------
diff --git a/dirty_cat/datasets/__init__.py b/dirty_cat/datasets/__init__.py
index 0507678..c011076 100644
--- a/dirty_cat/datasets/__init__.py
+++ b/dirty_cat/datasets/__init__.py
@@ -13,7 +13,7 @@ from ._fetching import (
     get_data_dir,
 )
 from ._generating import make_deduplication_data
-from ._ken_embeddings import get_ken_embeddings
+from ._ken_embeddings import get_ken_embeddings, get_ken_table_aliases, get_ken_types
 
 __all__ = [
     "DatasetAll",
@@ -30,4 +30,6 @@ __all__ = [
     "get_data_dir",
     "make_deduplication_data",
     "get_ken_embeddings",
+    "get_ken_table_aliases",
+    "get_ken_types",
 ]
diff --git a/dirty_cat/datasets/_ken_embeddings.py b/dirty_cat/datasets/_ken_embeddings.py
index cf592ae..c24ef8a 100644
--- a/dirty_cat/datasets/_ken_embeddings.py
+++ b/dirty_cat/datasets/_ken_embeddings.py
@@ -1,20 +1,117 @@
 """
 Get the Wikipedia embeddings for feature augmentation.
 """
+from typing import Optional, Set
+
 import pandas as pd
 from sklearn.decomposition import PCA
 
 from dirty_cat.datasets import fetch_figshare
 
+# Required for ignoring lines too long in the docstrings
+# flake8: noqa: E501
 
-def get_ken_embeddings(
-    types=None,
-    exclude=None,
-    embedding_table_id="all_entities",
-    embedding_type_id=None,
-    pca_components=None,
-    suffix="",
+
+_correspondence_table_url = (
+    "https://raw.githubusercontent.com/dirty-cat/datasets"
+    "/master/data/ken_correspondence.csv"
+)
+
+
+def get_ken_table_aliases() -> Set[str]:
+    """Get the supported aliases of embedded entities tables.
+
+    These aliases can be using in subsequent functions (see section *See Also*).
+
+    Returns
+    -------
+    aliases: set of str
+        The aliases of the embedded entities tables.
+
+    See Also
+    --------
+    :func:`get_ken_types`
+        Helper function to search for entity types.
+    :func:`get_ken_embeddings`
+        Download Wikipedia embeddings by type.
+
+    Notes
+    -----
+    Requires an Internet connection to work.
+    """
+    correspondence = pd.read_csv(_correspondence_table_url)
+    return set(["all_entities"] + list(correspondence["table"].values))
+
+
+def get_ken_types(
+    search: str = None,
+    *,
+    exclude: Optional[str] = None,
+    embedding_table_id: str = "all_entities",
 ):
+    """Helper function to search for entity types.
+
+    The result can then be used with :func:`get_ken_embeddings`.
+
+    Parameters
+    ----------
+    search : str, optional
+        Substring pattern that filters the types of entities.
+    exclude : str, optional
+        Substring pattern to exclude from the search.
+    embedding_table_id : str, optional, default='all_entities'
+        Table of embedded entities from which to extract the embeddings.
+        Get the supported tables with :func:`get_ken_table_aliases`.
+        It is NOT possible to pass a custom figshare ID.
+
+    Returns
+    -------
+    :class:`~pandas.DataFrame`
+        The types of entities containing the substring.
+
+    See Also
+    --------
+    :func:`get_ken_embeddings`
+        Download Wikipedia embeddings by type.
+
+    References
+    ----------
+    For more details, see Cvetkov-Iliev, A., Allauzen, A. & Varoquaux, G.:
+    `Relational data embeddings for feature enrichment
+    with background information. <https://doi.org/10.1007/s10994-022-06277-7>`_
+
+    Notes
+    -----
+    Best used in conjunction with :func:`get_ken_embeddings`.
+
+    """
+    correspondence = pd.read_csv(_correspondence_table_url)
+    if embedding_table_id not in get_ken_table_aliases():
+        raise ValueError(
+            f"The embedding_table_id must be one of {correspondence['table'].unique()}."
+        )
+    unique_types_figshare_id = correspondence[
+        correspondence["table"] == embedding_table_id
+    ]["unique_types_figshare_id"].values[0]
+    unique_types = fetch_figshare(unique_types_figshare_id)
+    if search is None:
+        search_result = unique_types.X
+    else:
+        search_result = unique_types.X[unique_types.X["Type"].str.contains(search)]
+    if exclude is not None:
+        search_result = search_result[~search_result["Type"].str.contains(exclude)]
+    return search_result
+
+
+def get_ken_embeddings(
+    types: Optional[str] = None,
+    *,
+    exclude: Optional[str] = None,
+    embedding_table_id: str = "all_entities",
+    embedding_type_id: Optional[str] = None,
+    pca_components: Optional[int] = None,
+    suffix: str = "",
+) -> pd.DataFrame:
     """Download Wikipedia embeddings by type.
 
     More details on the embeddings can be found on
@@ -22,36 +119,39 @@ def get_ken_embeddings(
 
     Parameters
     ----------
-    types : str, optional, default=None
+    types : str, optional
         Substring pattern that filters the types of entities.
         Will keep all entity types containing the substring.
         Write in lowercase. If None, all types will be passed.
-    exclude : str, optional, default=None
+    exclude : str, optional
         Type of embeddings to exclude from the types search.
-    embedding_table_id : {"all_entities", "albums", "companies", "movies", "games", "school"} or str, optional, default='all_entities' # noqa
+    embedding_table_id : str, optional, default='all_entities'
         Table of embedded entities from which to extract the embeddings.
-        See correspondence table
-        (https://github.com/dirty-cat/datasets/blob/master/data/ken_correspondence.csv)
-        for the figshare ID's of the tables.
-        It is also possible to introduce a custom figshare ID.
-    embedding_type_id : str, optional, default=None
-        Figshare ID of the file containing the type of embeddings. Ignored
-        unless a custom `embedding_table_id` is provided.
-    pca_components : int, optional, default=None
+        Get the supported tables with :func:`get_ken_table_aliases`.
+        It is also possible to pass a custom figshare ID.
+    embedding_type_id : str, optional
+        Figshare ID of the file containing the type of embeddings.
+        Get the supported tables with :func:`get_ken_types`.
+        Ignored unless a custom `embedding_table_id` is provided.
+    pca_components : int, optional
         Size of the dimensional space on which the embeddings will be projected
         by a principal component analysis.
-        If None, the default dimension (200) of the embeddins will be kept.
+        If None, the default dimension (200) of the embeddings will be kept.
     suffix : str, optional, default=""
         Suffix to add to the column names of the embeddings.
 
     Returns
     -------
-    embeddings: class:`~pandas.DataFrame`
+    :class:`~pandas.DataFrame`
         The embeddings of entities and the specified type from Wikipedia.
 
     See Also
     --------
-    :class:`~dirty_cat.fuzzy_join` :
+    :func:`get_ken_table_aliases`
+        Get the supported aliases of embedded entities tables.
+    :func:`get_ken_types`
+        Helper function to search for entity types.
+    :func:`~dirty_cat.fuzzy_join` :
         Join two tables (dataframes) based on approximate column matching.
     :class:`~dirty_cat.FeatureAugmenter` :
         Transformer to enrich a given table via one or more fuzzy joins to
@@ -75,21 +175,12 @@ def get_ken_embeddings(
     Going directly for the exact type name (e.g. "wikicat_rock_music_bands")
     is possible but may not be complete (as some relevant bands may be
     in other similar types).
-    For the full list of accepted types, see the `entity_detailed_types`
-    table on https://soda-inria.github.io/ken_embeddings/.
+    For searching the types, the :class:`~dirty_cat.datasets.get_ken_types`
+    function can be used.
 
     """
-    if embedding_table_id in [
-        "all_entities",
-        "albums",
-        "companies",
-        "movies",
-        "games",
-        "school",
-    ]:
-        correspondence = pd.read_csv(
-            "https://raw.githubusercontent.com/dirty-cat/datasets/master/data/ken_correspondence.csv"  # noqa
-        )
+    if embedding_table_id in get_ken_table_aliases():
+        correspondence = pd.read_csv(_correspondence_table_url)
         embeddings_id = correspondence[correspondence["table"] == embedding_table_id][
             "entities_figshare_id"
         ].values[0]
@@ -105,7 +196,6 @@ def get_ken_embeddings(
         emb_type = emb_type[~emb_type["Type"].str.contains(exclude)]
     emb_type.drop_duplicates(subset=["Entity"], inplace=True)
     emb_final = []
-    emb_df = pd.DataFrame()
     emb_full = fetch_figshare(embeddings_id)
     for path in emb_full.path:
         emb_extracts = pd.read_parquet(path)
diff --git a/examples/06_ken_embeddings_example.py b/examples/06_ken_embeddings_example.py
index 1e23bcb..9502ead 100644
--- a/examples/06_ken_embeddings_example.py
+++ b/examples/06_ken_embeddings_example.py
@@ -9,7 +9,7 @@ sources may be the key to improving the analysis.
 Embeddings, or vectorial representations of entities, are a conveniant way to
 capture and summarize the information on an entity.
 Relational data embeddings capture all common entities from Wikipedia. [#]_
-These will be called `KEN Embeddings` in the following example.
+These will be called `KEN embeddings` in the following example.
 
 We will see that these embeddings of common entities significantly
 improve our results.
@@ -29,9 +29,6 @@ improve our results.
  .. |MinHash| replace::
      :class:`~dirty_cat.MinHashEncoder`
 
- .. |KEN| replace::
-     :class:`~dirty_cat.datasets.get_ken_embeddings`
-
  .. |HGBR| replace::
      :class:`~sklearn.ensemble.HistGradientBoostingRegressor`
 """
@@ -40,8 +37,8 @@ improve our results.
 # The data
 # --------
 #
-# We will take a look at the Video Game sales dataset.
-# We first retrieve the dataset:
+# We will take a look at the video game sales dataset.
+# Let's retrieve the dataset:
 import pandas as pd
 
 X = pd.read_csv(
@@ -49,12 +46,12 @@ X = pd.read_csv(
     sep=";",
     on_bad_lines="skip",
 )
-# We shuffle the data
+# Shuffle the data
 X = X.sample(frac=1, random_state=11, ignore_index=True)
 X.head(3)
 
 ###############################################################################
-# Our goal will be to predict y, our target column (the sales amount):
+# Our goal will be to predict the sales amount (y, our target column):
 y = X["Global_Sales"]
 y
 
@@ -69,7 +66,7 @@ sns.histplot(y)
 plt.show()
 
 ###############################################################################
-# It seems better to take the log of sales rather than the sales amount itself:
+# It seems better to take the log of sales rather than the absolute values:
 import numpy as np
 
 y = np.log(y)
@@ -77,25 +74,44 @@ sns.histplot(y)
 plt.show()
 
 ###############################################################################
-# Now, let's carry out some basic preprocessing:
+# Before moving further, let's carry out some basic preprocessing:
 
 # Get a mask of the rows with missing values in "Publisher" and "Global_Sales"
-mask1 = X.isna()["Publisher"]
-mask2 = X.isna()["Global_Sales"]
+mask = X.isna()["Publisher"] | X.isna()["Global_Sales"]
 # And remove them
 X.dropna(subset=["Publisher", "Global_Sales"], inplace=True)
-y = y[~mask1]
-y = y[~mask2]
+y = y[~mask]
 
 ###############################################################################
 # Extracting entity embeddings
 # ----------------------------
-# We will use the |KEN| function to extract the embeddings
-# of entities we need:
+#
+# We will use KEN embeddings to enrich our data.
+#
+# We will start by checking out the available tables with
+# :class:`~dirty_cat.datasets.get_ken_table_aliases`:
+from dirty_cat.datasets import get_ken_table_aliases
+
+get_ken_table_aliases()
+
+###############################################################################
+# The *games* table is the most relevant to our case.
+# Let's see what kind of types we can find in it with the function
+# :class:`~dirty_cat.datasets.get_ken_types`:
+from dirty_cat.datasets import get_ken_types
+
+get_ken_types(embedding_table_id="games")
+
+###############################################################################
+# Interesting, we have a broad range of topics!
+#
+# Next, we'll use :class:`~dirty_cat.datasets.get_ken_embeddings`
+# to extract the embeddings of entities we need:
 from dirty_cat.datasets import get_ken_embeddings
 
 ###############################################################################
-# KEN Embeddings are classified by types. The |KEN| function
+# KEN Embeddings are classified by types.
+# The :class:`~dirty_cat.datasets.get_ken_embeddings` function
 # allows us to specify the types to be included and/or excluded
 # so as not to load all Wikipedia entity embeddings in a table.
 #


Test patch:
diff --git a/dirty_cat/datasets/tests/test_ken_embeddings.py b/dirty_cat/datasets/tests/test_ken_embeddings.py
index 125789f..c30e004 100644
--- a/dirty_cat/datasets/tests/test_ken_embeddings.py
+++ b/dirty_cat/datasets/tests/test_ken_embeddings.py
@@ -1,4 +1,29 @@
-from dirty_cat.datasets import get_ken_embeddings
+from dirty_cat.datasets import get_ken_embeddings, get_ken_table_aliases, get_ken_types
+
+
+def test_get_ken_table_aliases():
+    """
+    Test if the aliases of the tables are correctly fetched
+    """
+    aliases = get_ken_table_aliases()
+    assert "all_entities" in aliases
+    assert "games" in aliases
+    assert "albums" in aliases
+
+
+def test_get_ken_types():
+    """
+    Test if the types of entities are correctly fetched
+    """
+    # Tests the full result returns alright
+    types = get_ken_types()
+    assert types.shape[0] == 114509
+    # Tests the additive filter works
+    types2 = get_ken_types(search="game")
+    assert types2.shape[0] == 2540
+    # Tests the negative filter works
+    types3 = get_ken_types(search="game", exclude="card")
+    assert types3.shape[0] == 2532
 
 
 def test_small_ken_embeddings():