assert datasettemplates type

michaelbyun · Apr 12, 2023 · a7f5a8b · a7f5a8b
1 parent 6d46f7e
commit a7f5a8b
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/elk/extraction/prompt_loading.py b/elk/extraction/prompt_loading.py
@@ -121,8 +121,9 @@ def __post_init__(self):
 
  # Write to file if successfully merged all prompts.
  if self.combined_template_path:
- print("Total number of templates: ", len(combined_prompter.templates))
- combined_prompter.write_to_file()
+ prompter = assert_type(DatasetTemplates, combined_prompter)
+ print("Total number of templates: ", len(prompter.templates))
+ prompter.write_to_file()
  print(
  "Saved to promptsource/templates/combined_templates/"
  + f"{self.combined_template_path}.yaml"