fix mypy error

sotopia-lab · Jan 25, 2024 · d979d67 · d979d67
1 parent 251a773
commit d979d67
Showing 1 changed file with 5 additions and 6 deletions.
diff --git a/examples/fix_missing_episodes_with_tag.py b/examples/fix_missing_episodes_with_tag.py
@@ -74,12 +74,11 @@ def get_all_episodes(tags: List[str] = []) -> List[Tuple[EpisodeLog, str]]:
     assert len(tags) != 0, "No episodes found"
     if len(tags) != 0:
         print("Using tag list: ", tags)
-        episode_pks = []
+        episode_pks: Tuple[(str, str)] = []
         for tag in tags:
-            episode_pks += [
-                (ep.pk, ep.tag)
-                for ep in EpisodeLog.find(EpisodeLog.tag == tag).all()
-            ]
+            ep_list = list(EpisodeLog.find(EpisodeLog.tag == tag).all())
+            ep_list = cast(List[EpisodeLog], ep_list)
+            episode_pks += [(ep.pk, ep.tag) for ep in ep_list]
 
     all_episodes = []
     for pk, tag in tqdm(episode_pks):
@@ -135,7 +134,7 @@ def find_combo_pk(
 def get_combo_model_map(
     all_episodes: List[Tuple[EpisodeLog, str]],
     all_combos_map: Dict[str, EnvAgentComboStorage],
-):
+) -> Dict[str, Counter[tuple[LLM_Name, LLM_Name, LLM_Name, str]]]:
     combo_model_map: Dict[
         str, Counter[tuple[LLM_Name, LLM_Name, LLM_Name, str]]
     ] = defaultdict(Counter)