Replace tensorflow ST5 embedder with torch based one

This will ensure the embedder works with M1-based macs (fixes #37). PiperOrigin-RevId: 625611229 Change-Id: I6591565e82774e23aa40bd7a33571f8fe6a9a0d5
google-deepmind · Apr 17, 2024 · d0df6e3 · d0df6e3
1 parent 449cd8f
commit d0df6e3
Show file tree

Hide file tree

Showing 9 changed files with 36 additions and 85 deletions.
diff --git a/concordia/associative_memory/embedder_st5.py b/concordia/associative_memory/embedder_st5.py
diff --git a/examples/cyberball/cyberball.ipynb b/examples/cyberball/cyberball.ipynb
@@ -71,6 +71,7 @@
         "\n",
         "from google.colab import widgets  # pytype: disable=import-error\n",
         "from IPython import display\n",
+        "import sentence_transformers\n",
         "\n",
         "from concordia.agents import basic_agent\n",
         "from concordia import components as generic_components\n",
@@ -79,7 +80,6 @@
         "from concordia.document import interactive_document\n",
         "from concordia.associative_memory import associative_memory\n",
         "from concordia.associative_memory import blank_memories\n",
-        "from concordia.associative_memory import embedder_st5\n",
         "from concordia.associative_memory import formative_memories\n",
         "from concordia.associative_memory import importance_function\n",
         "from concordia.clocks import game_clock\n",
@@ -106,7 +106,8 @@
       "outputs": [],
       "source": [
         "# Setup sentence encoder\n",
-        "embedder = embedder_st5.EmbedderST5()"
+        "st_model = sentence_transformers.SentenceTransformer('sentence-transformers/all-mpnet-base-v2')\n",
+        "embedder = lambda x: st_model.encode(x, show_progress_bar=False)\n"
       ]
     },
     {

diff --git a/examples/magic_beans_for_sale.ipynb b/examples/magic_beans_for_sale.ipynb
@@ -68,6 +68,7 @@
         "import datetime\n",
         "\n",
         "import numpy as np\n",
+        "import sentence_transformers\n",
         "\n",
         "from google.colab import widgets  # pytype: disable=import-error\n",
         "from IPython import display\n",
@@ -77,7 +78,6 @@
         "from concordia import components as generic_components\n",
         "from concordia.associative_memory import associative_memory\n",
         "from concordia.associative_memory import blank_memories\n",
-        "from concordia.associative_memory import embedder_st5\n",
         "from concordia.associative_memory import formative_memories\n",
         "from concordia.associative_memory import importance_function\n",
         "from concordia.clocks import game_clock\n",
@@ -102,7 +102,8 @@
       "outputs": [],
       "source": [
         "# Setup sentence encoder\n",
-        "embedder = embedder_st5.EmbedderST5()"
+        "st_model = sentence_transformers.SentenceTransformer('sentence-transformers/all-mpnet-base-v2')\n",
+        "embedder = lambda x: st_model.encode(x, show_progress_bar=False)\n"
       ]
     },
     {

diff --git a/examples/phone/calendar.ipynb b/examples/phone/calendar.ipynb
@@ -69,14 +69,14 @@
         "import random\n",
         "\n",
         "from IPython import display\n",
+        "import sentence_transformers\n",
         "\n",
         "from concordia import components as generic_components\n",
         "from concordia.components import agent as components\n",
         "from concordia.components import game_master as gm_components\n",
         "from concordia.agents import basic_agent\n",
         "from concordia.associative_memory import associative_memory\n",
         "from concordia.associative_memory import blank_memories\n",
-        "from concordia.associative_memory import embedder_st5\n",
         "from concordia.associative_memory import formative_memories\n",
         "from concordia.associative_memory import importance_function\n",
         "from concordia.clocks import game_clock\n",
@@ -99,7 +99,8 @@
       "outputs": [],
       "source": [
         "#@title Setup sentence encoder\n",
-        "embedder = embedder_st5.EmbedderST5()"
+        "st_model = sentence_transformers.SentenceTransformer('sentence-transformers/all-mpnet-base-v2')\n",
+        "embedder = lambda x: st_model.encode(x, show_progress_bar=False)\n"
       ]
     },
     {

diff --git a/examples/requirements.txt b/examples/requirements.txt
@@ -1,4 +1,5 @@
 absl-py
 docstring-parser
 gdm-concordia
+sentence_transformers
 termcolor
diff --git a/examples/three_key_questions.ipynb b/examples/three_key_questions.ipynb
@@ -75,6 +75,7 @@
         "import datetime\n",
         "\n",
         "import numpy as np\n",
+        "import sentence_transformers\n",
         "\n",
         "from google.colab import widgets  # pytype: disable=import-error\n",
         "from IPython import display\n",
@@ -84,7 +85,6 @@
         "from concordia import components as generic_components\n",
         "from concordia.associative_memory import associative_memory\n",
         "from concordia.associative_memory import blank_memories\n",
-        "from concordia.associative_memory import embedder_st5\n",
         "from concordia.associative_memory import formative_memories\n",
         "from concordia.associative_memory import importance_function\n",
         "from concordia.clocks import game_clock\n",
@@ -109,7 +109,8 @@
       "outputs": [],
       "source": [
         "# Setup sentence encoder\n",
-        "embedder = embedder_st5.EmbedderST5()"
+        "st_model = sentence_transformers.SentenceTransformer('sentence-transformers/all-mpnet-base-v2')\n",
+        "embedder = lambda x: st_model.encode(x, show_progress_bar=False)\n"
       ]
     },
     {

diff --git a/examples/village/day_in_riverbend.ipynb b/examples/village/day_in_riverbend.ipynb
@@ -73,14 +73,14 @@
         "\n",
         "from google.colab import widgets  # pytype: disable=import-error\n",
         "from IPython import display\n",
+        "import sentence_transformers\n",
         "\n",
         "from concordia import components as generic_components\n",
         "from concordia.agents import basic_agent\n",
         "from concordia.components import agent as components\n",
         "from concordia.agents import basic_agent\n",
         "from concordia.associative_memory import associative_memory\n",
         "from concordia.associative_memory import blank_memories\n",
-        "from concordia.associative_memory import embedder_st5\n",
         "from concordia.associative_memory import formative_memories\n",
         "from concordia.associative_memory import importance_function\n",
         "from concordia.clocks import game_clock\n",
@@ -105,7 +105,8 @@
       "outputs": [],
       "source": [
         "# @title Setup sentence encoder\n",
-        "embedder = embedder_st5.EmbedderST5()"
+        "st_model = sentence_transformers.SentenceTransformer('sentence-transformers/all-mpnet-base-v2')\n",
+        "embedder = lambda x: st_model.encode(x, show_progress_bar=False)\n"
       ]
     },
     {

diff --git a/examples/village/riverbend_elections.ipynb b/examples/village/riverbend_elections.ipynb
@@ -71,14 +71,14 @@
         "\n",
         "from google.colab import widgets  # pytype: disable=import-error\n",
         "from IPython import display\n",
+        "import sentence_transformers\n",
         "\n",
         "from concordia import components as generic_components\n",
         "from concordia.components import agent as components\n",
         "from concordia.components import game_master as gm_components\n",
         "from concordia.agents import basic_agent\n",
         "from concordia.associative_memory import associative_memory\n",
         "from concordia.associative_memory import blank_memories\n",
-        "from concordia.associative_memory import embedder_st5\n",
         "from concordia.associative_memory import formative_memories\n",
         "from concordia.associative_memory import importance_function\n",
         "from concordia.clocks import game_clock\n",
@@ -104,7 +104,8 @@
       "outputs": [],
       "source": [
         "# @title Setup sentence encoder\n",
-        "embedder = embedder_st5.EmbedderST5()"
+        "st_model = sentence_transformers.SentenceTransformer('sentence-transformers/all-mpnet-base-v2')\n",
+        "embedder = lambda x: st_model.encode(x, show_progress_bar=False)\n"
       ]
     },
     {

diff --git a/setup.py b/setup.py
@@ -15,34 +15,8 @@
 
 """Install script for setuptools."""
 
-import platform
 import setuptools
 
-IS_M1_OSX = platform.system() == 'Darwin' and platform.machine() == 'arm64'
-
-REQUIREMENTS = (
-    # TODO: b/312199199 - remove some requirements.
-    'absl-py',
-    'google-cloud-aiplatform',
-    'ipython',
-    'langchain',
-    'matplotlib',
-    'numpy',
-    'openai>=1.3.0',
-    'pandas<=2.0.3',
-    'python-dateutil',
-    'reactivex',
-    'retry',
-    'scipy',
-    'tensorflow',
-    'tensorflow-hub',
-    'tensorflow-text',
-    'termcolor',
-    'typing-extensions',
-)
-M1_OSX_REQUIREMENTS = tuple(set(REQUIREMENTS) - {'tensorflow-text'})
-
-
 setuptools.setup(
     name='gdm-concordia',
     version='1.2.0',
@@ -75,7 +49,23 @@
     packages=setuptools.find_packages(include=['concordia', 'concordia.*']),
     package_data={},
     python_requires='>=3.11',
-    install_requires=M1_OSX_REQUIREMENTS if IS_M1_OSX else REQUIREMENTS,
+    install_requires=(
+        # TODO: b/312199199 - remove some requirements.
+        'absl-py',
+        'google-cloud-aiplatform',
+        'ipython',
+        'langchain',
+        'matplotlib',
+        'numpy',
+        'openai>=1.3.0',
+        'pandas<=2.0.3',
+        'python-dateutil',
+        'reactivex',
+        'retry',
+        'scipy',
+        'termcolor',
+        'typing-extensions',
+    ),
     extras_require={
         # Used in development.
         'dev': [