Fixing imports and colab/codespace

PiperOrigin-RevId: 661180043 Change-Id: I0578ca6ce1734b7e8d8f879d8104fbfcff305ada
google-deepmind · Aug 9, 2024 · dba1105 · dba1105
1 parent a8c9d67
commit dba1105
Showing 1 changed file with 25 additions and 15 deletions.
diff --git a/examples/modular/notebook.ipynb b/examples/modular/notebook.ipynb
@@ -38,18 +38,24 @@
       },
       "outputs": [],
       "source": [
-        "!pip install --ignore-requires-python git+https://github.com/google-deepmind/concordia.git"
-      ]
-    },
-    {
-      "cell_type": "code",
-      "execution_count": null,
-      "metadata": {
-        "id": "B_JQ_dJDpBVC"
-      },
-      "outputs": [],
-      "source": [
-        "!pip install -r https://raw.githubusercontent.com/google-deepmind/concordia/main/examples/requirements.txt"
+        "!pip install -r https://raw.githubusercontent.com/google-deepmind/concordia/main/examples/requirements.txt\n",
+        "!pip install --ignore-requires-python git+https://github.com/google-deepmind/concordia.git\n",
+        "\n",
+        "#@markdown Here we check if this is being ran in colab or codespaces. If running a colab, we clone the repository to get the example environments\n",
+        "import os\n",
+        "import sys\n",
+        "import pathlib\n",
+        "\n",
+        "if os.environ.get('CODESPACES') == 'true':\n",
+        "  print(\"Running in Codespaces\")\n",
+        "  concordia_root_dir = pathlib.Path.cwd().parent.parent.resolve()\n",
+        "  sys.path.append(f'{concordia_root_dir}')\n",
+        "elif 'google.colab' in sys.modules:\n",
+        "  print(\"Running in Colab\")\n",
+        "  !git clone https://github.com/google-deepmind/concordia.git\n",
+        "  sys.path.append('/content/concordia/')\n",
+        "else:\n",
+        "  print(\"Running in a different environment, please make sure to add root directory of examples to with sys.path.append\")"
       ]
     },
     {
@@ -66,10 +72,16 @@
         "import pathlib\n",
         "import sys\n",
         "\n",
+        "from concordia.language_model import amazon_bedrock_model\n",
+        "from concordia.language_model import google_aistudio_model\n",
         "from concordia.language_model import gpt_model\n",
+        "from concordia.language_model import langchain_ollama_model\n",
         "from concordia.language_model import mistral_model\n",
         "from concordia.language_model import no_language_model\n",
+        "from concordia.language_model import ollama_model\n",
+        "from concordia.language_model import pytorch_gemma_model\n",
         "from concordia.utils import measurements as measurements_lib\n",
+        "\n",
         "import openai\n",
         "import sentence_transformers"
       ]
@@ -134,9 +146,7 @@
         "agent_module = importlib.import_module(\n",
         "    f'{IMPORT_AGENT_BASE_DIR}.{AGENT_NAME}')\n",
         "# Load the environment config with importlib\n",
-        "concordia_root_dir = pathlib.Path.cwd().parent.parent.resolve()\n",
-        "sys.path.append(f'{concordia_root_dir}')\n",
-        "IMPORT_ENV_BASE_DIR = 'environment'\n",
+        "IMPORT_ENV_BASE_DIR = 'examples.modular.environment'\n",
         "simulation = importlib.import_module(\n",
         "    f'{IMPORT_ENV_BASE_DIR}.{ENVIRONMENT_NAME}')"
       ]