Refactor imports and update type annotations in redis_stats.ipynb

sotopia-lab · Apr 4, 2024 · 00a12e3 · 00a12e3
1 parent f9f15a3
commit 00a12e3
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 5 deletions.
diff --git a/notebooks/redis_stats.ipynb b/notebooks/redis_stats.ipynb
@@ -10,7 +10,6 @@
     "import os\n",
     "import json\n",
     "from typing import get_args\n",
-    "from sotopia.generation_utils.generate import LLM_Name\n",
     "from tqdm.notebook import tqdm\n",
     "import rich\n",
     "import logging\n",
@@ -37,6 +36,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "LLM_Name = Literal[\n",
+    "    \"togethercomputer/llama-2-7b-chat\",\n",
+    "    \"togethercomputer/llama-2-70b-chat\",\n",
+    "    \"togethercomputer/mpt-30b-chat\",\n",
+    "    \"gpt-3.5-turbo\",\n",
+    "    \"text-davinci-003\",\n",
+    "    \"gpt-4\",\n",
+    "    \"gpt-4-turbo\",\n",
+    "    \"human\",\n",
+    "    \"redis\",\n",
+    "]\n",
     "def _is_valid_episode_log_pk(pk: str) -> bool:\n",
     "    try:\n",
     "        episode = EpisodeLog.get(pk=pk)\n",

diff --git a/sotopia/database/utils.py b/sotopia/database/utils.py
@@ -11,11 +11,11 @@
 
 class TwoAgentEpisodeWithScenarioBackgroundGoals(JsonModel):
     episode_id: str = Field(index=True)
-    scenario: str = Field(index=True)
+    scenario: str = Field(index=False)
     codename: str = Field(index=True)
-    agents_background: dict[str, str] = Field(index=True)
-    social_goals: dict[str, str] = Field(index=True)
-    social_interactions: str = Field(index=True)
+    agents_background: dict[str, str] = Field(index=False)
+    social_goals: dict[str, str] = Field(index=False)
+    social_interactions: str = Field(index=False)
 
 
 def _map_gender_to_adj(gender: str) -> str: