import logging
from dataclasses import dataclass
from enum import Enum
from typing import Any, Optional
from ... import list_step_functions
from ...attr.step_functions import is_contrastive_step_function
from ...utils import cli_arg, pretty_dict
from ..attribute import AttributeBaseArgs
from ..commands_utils import command_args_docstring
logger = logging.getLogger(__name__)
class HandleOutputContextSetting(Enum):
MANUAL = "manual"
AUTO = "auto"
PRE = "pre"
@command_args_docstring
@dataclass
class AttributeContextInputArgs:
input_current_text: str = cli_arg(
default="",
help=(
"The input text used for generation. If the model is a decoder-only model, the input text is a "
"prompt used for language modeling. If the model is an encoder-decoder model, the input text is the "
"source text provided as input to the encoder. It will be formatted as {current} in the "
"``input_template``."
),
)
input_context_text: Optional[str] = cli_arg(
default=None,
help=(
"Additional input context influencing the generation of ``output_current_text``. If the model is a"
" decoder-only model, the input text is a prefix to the ``input_current_text`` prompt. If the model is an"
" encoder-decoder model, the input context is part of the source text provided as input to the encoder. "
" It will be formatted as {context} in the ``input_template``."
),
)
input_template: Optional[str] = cli_arg(
default=None,
help=(
"The template used to format model inputs. The template must contain at least the"
" ``{current}`` placeholder, which will be replaced by ``input_current_text``. If ``{context}`` is"
" also specified, input-side context will be used. Can be modified for models requiring special tokens or"
" formatting in the input text (e.g. <brk> tags to separate context and current inputs)."
" Defaults to '{context} {current}' if ``input_context_text`` is provided, '{current}' otherwise."
),
)
output_context_text: Optional[str] = cli_arg(
default=None,
help=(
"An output contexts for which context sensitivity should be detected. For encoder-decoder models, this"
" is a target-side prefix to the output_current_text used as input to the decoder. For decoder-only "
" models this is a portion of the model generation that should be considered as an additional context "
" (e.g. a chain-of-thought sequence). It will be formatted as {context} in the ``output_template``."
" If not provided but specified in the ``output_template``, the output context will be generated"
" along with the output current text, and user validation might be required to separate the two."
),
)
output_current_text: Optional[str] = cli_arg(
default=None,
help=(
"The output text generated by the model when all available contexts are provided. Tokens in "
" ``output_current_text`` will be tested for context-sensitivity, and their generation will be attributed "
" to input/target contexts (if present) in case they are found to be context-sensitive. If specified, "
" this output is force-decoded. Otherwise, it is generated by the model using infilled ``input_template`` "
" and ``output_template``. It will be formatted as {current} in the ``output_template``."
),
)
output_template: Optional[str] = cli_arg(
default=None,
help=(
"The template used to format model outputs. The template must contain at least the"
" ``{current}`` placeholder, which will be replaced by ``output_current_text``. If ``{context}`` is"
" also specified, output-side context will be used. Can be modified for models requiring special tokens or"
" formatting in the output text (e.g. <brk> tags to separate context and current outputs)."
" Defaults to '{context} {current}' if ``output_context_text`` is provided, '{current}' otherwise."
),
)
contextless_input_current_text: Optional[str] = cli_arg(
default=None,
help=(
"The input current text or template to use in the contrastive comparison with contextual input. By default"
" it is the same as ``input_current_text``, but it can be useful in cases where the context is nested "
"inside the current text (e.g. for an ``input_template`` like <user>\n{context}\n{current}\n<assistant> we "
"can use this parameter to format the contextless version as <user>\n{current}\n<assistant>)."
"If it contains the tag {current}, it will be infilled with the ``input_current_text``. Otherwise, it will"
" be used as-is for the contrastive comparison, enabling contrastive comparison with different inputs."
),
)
contextless_output_current_text: Optional[str] = cli_arg(
default=None,
help=(
"The output current text or template to use in the contrastive comparison with contextual output. By default"
" it is the same as ``output_current_text``, but it can be useful in cases where the context is nested "
"inside the current text (e.g. for an ``output_template`` like <user>\n{context}\n{current}\n<assistant> we "
"can use this parameter to format the contextless version as <user>\n{current}\n<assistant>)."
"If it contains the tag {current}, it will be infilled with the ``output_current_text``. Otherwise, it will"
" be used as-is for the contrastive comparison, enabling contrastive comparison with different outputs."
),
)
@command_args_docstring
@dataclass
class AttributeContextMethodArgs(AttributeBaseArgs):
context_sensitivity_metric: str = cli_arg(
default="kl_divergence",
help="The contrastive metric used to detect context-sensitive tokens in ``output_current_text``.",
choices=[fn for fn in list_step_functions() if is_contrastive_step_function(fn)],
)
handle_output_context_strategy: str = cli_arg(
default=HandleOutputContextSetting.MANUAL.value,
choices=[e.value for e in HandleOutputContextSetting],
help=(
"Specifies how output context should be handled when it is produced together with the output current text,"
" and the two need to be separated for context sensitivity detection.\n"
"Options:\n"
"- `manual`: The user is prompted to verify an automatic context detection attempt, and optionally to"
" provide a correct context separation manually.\n"
"- `auto`: Attempts an automatic detection of context using an automatic alignment with source context"
" (assuming an MT-like task).\n"
"- `pre`: If context is required but not pre-defined by the user via the ``output_context_text`` argument,"
" the execution will fail instead of attempting to prompt the user for the output context."
),
)
contextless_output_next_tokens: list[str] = cli_arg(
default_factory=list,
help=(
"If specified, it should provide a list of one token per CCI output indicating the next token that should"
" be force-decoded as contextless output instead of the natural output produced by"
" ``get_contextless_output``. This is ignored if the ``attributed_fn`` used is not contrastive."
),
)
prompt_user_for_contextless_output_next_tokens: bool = cli_arg(
default=False,
help=(
"If specified, the user is prompted to provide the next token that should be force-decoded as contextless"
" output instead of the natural output produced by ``get_contextless_output``. This is ignored if the"
" ``attributed_fn`` used is not contrastive."
),
)
special_tokens_to_keep: list[str] = cli_arg(
default_factory=list,
help="Special tokens to preserve in the generated string, e.g. ``<brk>`` separator between context and current.",
)
decoder_input_output_separator: str = cli_arg(
default=" ",
help=(
"If specified, the separator used to split the input and output of the decoder. If not specified, the"
" separator is a whitespace character."
),
)
context_sensitivity_std_threshold: float = cli_arg(
default=1.0,
help=(
"Parameter to control the selection of ``output_current_text`` tokens considered as context-sensitive for "
"moving onwards with attribution. Corresponds to the number of standard deviations above or below the mean"
" ``context_sensitivity_metric`` score for tokens to be considered context-sensitive."
),
)
context_sensitivity_topk: Optional[int] = cli_arg(
default=None,
help=(
"If set, after selecting the salient context-sensitive tokens with ``context_sensitivity_std_threshold`` "
"only the top-K remaining tokens are used. By default no top-k selection is performed."
),
)
attribution_std_threshold: float = cli_arg(
default=1.0,
help=(
"Parameter to control the selection of ``input_context_text`` and ``output_context_text`` tokens "
"considered as salient as a result for the attribution process. Corresponds to the number of standard "
"deviations above or below the mean ``attribution_method`` score for tokens to be considered salient. "
"CCI scores for all context tokens are saved in the output, but this parameter controls which tokens are "
"used in the visualization of context reliance."
),
)
attribution_topk: Optional[int] = cli_arg(
default=None,
help=(
"If set, after selecting the most salient tokens with ``attribution_std_threshold`` "
"only the top-K remaining tokens are used. By default no top-k selection is performed."
),
)
@command_args_docstring
@dataclass
class AttributeContextOutputArgs:
show_intermediate_outputs: bool = cli_arg(
default=False,
help=(
"If specified, the intermediate outputs produced by the Inseq library for context-sensitive target "
"identification (CTI) and contextual cues imputation (CCI) are shown during the process."
),
)
save_path: Optional[str] = cli_arg(
default=None,
aliases=["-o"],
help="If present, the output of the two-step process will be saved in JSON format at the specified path.",
)
add_output_info: bool = cli_arg(
default=True,
help="If specified, additional information about the attribution process is added to the saved output.",
)
viz_path: Optional[str] = cli_arg(
default=None,
help="If specified, the visualization produced from the output is saved in HTML format at the specified path.",
)
show_viz: bool = cli_arg(
default=True,
help="If specified, the visualization produced from the output is shown in the terminal.",
)
[docs]
@command_args_docstring
@dataclass
class AttributeContextArgs(AttributeContextInputArgs, AttributeContextMethodArgs, AttributeContextOutputArgs):
def __repr__(self):
return f"{self.__class__.__name__}({pretty_dict(self.__dict__)})"
@classmethod
def _to_dict(cls, val: Any) -> dict[str, Any]:
if val is None or isinstance(val, (str, int, float, bool)):
return val
elif isinstance(val, dict):
return {k: cls._to_dict(v) for k, v in val.items()}
elif isinstance(val, (list, tuple)):
return [cls._to_dict(v) for v in val]
else:
return str(val)
def to_dict(self) -> dict[str, Any]:
return self._to_dict(self.__dict__)
def __post_init__(self):
if (
self.handle_output_context_strategy == HandleOutputContextSetting.PRE.value
and not self.output_context_text
and "{context}" in self.output_template
):
raise ValueError(
"If --handle_output_context_strategy='pre' and {context} is used in --output_template, --output_context_text"
" must be specified to avoid user prompt for output context."
)
if len(self.contextless_output_next_tokens) > 0 and self.prompt_user_for_contextless_output_next_tokens:
raise ValueError(
"Only one of contextless_output_next_tokens and prompt_user_for_contextless_output_next_tokens can be"
" specified."
)
if self.input_template is None:
self.input_template = "{current}" if self.input_context_text is None else "{context} {current}"
if self.output_template is None:
self.output_template = "{current}" if self.output_context_text is None else "{context} {current}"
if self.contextless_input_current_text is None:
self.contextless_input_current_text = "{current}"
if "{current}" not in self.contextless_input_current_text:
raise ValueError(
"{current} placeholder is missing from contextless_input_current_text template"
f" {self.contextless_input_current_text}."
)
if self.contextless_output_current_text is None:
self.contextless_output_current_text = "{current}"
if "{current}" not in self.contextless_output_current_text:
raise ValueError(
"{current} placeholder is missing from contextless_output_current_text template"
f" {self.contextless_output_current_text}."
)
self.has_input_context = "{context}" in self.input_template
self.has_output_context = "{context}" in self.output_template
if not self.input_current_text:
raise ValueError("--input_current_text must be a non-empty string.")
if self.input_context_text and not self.has_input_context:
logger.warning(
f"input_template has format {self.input_template} (no {{context}}), but --input_context_text is"
" specified. Ignoring provided --input_context_text."
)
self.input_context_text = None
if self.output_context_text and not self.has_output_context:
logger.warning(
f"output_template has format {self.output_template} (no {{context}}), but --output_context_text is"
" specified. Ignoring provided --output_context_text."
)
self.output_context_text = None
if not self.input_context_text and self.has_input_context:
raise ValueError(
f"{{context}} format placeholder is present in input_template {self.input_template},"
" but --input_context_text is not specified."
)
if "{current}" not in self.input_template:
raise ValueError(f"{{current}} format placeholder is missing from input_template {self.input_template}.")
if "{current}" not in self.output_template:
raise ValueError(f"{{current}} format placeholder is missing from output_template {self.output_template}.")
if not self.input_current_text:
raise ValueError("--input_current_text must be a non-empty string.")
if self.has_output_context and self.output_template.find("{context}") > self.output_template.find("{current}"):
raise ValueError(
f"{{context}} placeholder must appear before {{current}} in output_template '{self.output_template}'."
)
if not self.output_template.endswith("{current}"):
*_, suffix = self.output_template.partition("{current}")
logger.warning(
f"Suffix '{suffix}' was specified in output_template and will be used to ignore the specified suffix"
" tokens during context sensitivity detection. Make sure that the suffix corresponds to the end of the"
" output_current_text by forcing --output_current_text if necessary."
)