fix: register placeholder_to_record_key in GeneratePrompt transform (#…

…301)
aws · Jul 2, 2024 · a911e56 · a911e56
1 parent 00033e9
commit a911e56
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/fmeval/transforms/common.py b/src/fmeval/transforms/common.py
@@ -46,7 +46,8 @@ def __init__(
         {"my_prompt": "Summarize some long text and some other long text"}
         """
         super().__init__(input_keys, output_keys, prompt_template, placeholder_to_record_key)
-        self.register_input_output_keys(input_keys, output_keys)
+        input_keys_to_register = list(placeholder_to_record_key.values()) if placeholder_to_record_key else input_keys
+        self.register_input_output_keys(input_keys_to_register, output_keys)
         self.placeholder_to_record_key = placeholder_to_record_key
         self.prompt_template = prompt_template
         self.prompt_composer = PromptComposer(prompt_template)