Permit none key in chain mapper (#8696)

This commit is contained in:
William FH 2023-08-03 08:50:36 -07:00 committed by GitHub
parent e283dc8d50
commit a137492b53
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -126,9 +126,12 @@ class ChainStringRunMapper(StringRunMapper):
"""Extract items to evaluate from the run object from a chain."""
input_key: Optional[str] = None
"""The key from the model Run's inputs to use as the eval input."""
"""The key from the model Run's inputs to use as the eval input.
If not provided, will use the only input key or raise an error if there are multiple.
"""
prediction_key: Optional[str] = None
"""The key from the model Run's outputs to use as the eval prediction."""
"""The key from the model Run's outputs to use as the eval prediction.
If not provided, will use the only output key or raise an error if there are multiple."""
def _get_key(self, source: Dict, key: Optional[str], which: str) -> str:
if key is not None:
@ -145,11 +148,9 @@ class ChainStringRunMapper(StringRunMapper):
"""Maps the Run to a dictionary."""
if not run.outputs:
raise ValueError(f"Run {run.id} has no outputs to evaluate.")
if run.run_type != "chain":
raise ValueError("Chain RunMapper only supports Chain runs.")
if self.input_key not in run.inputs:
if self.input_key is not None and self.input_key not in run.inputs:
raise ValueError(f"Run {run.id} does not have input key {self.input_key}.")
elif self.prediction_key not in run.outputs:
elif self.prediction_key is not None and self.prediction_key not in run.outputs:
raise ValueError(
f"Run {run.id} does not have prediction key {self.prediction_key}."
)