| from typing import Optional | |
| from inseq.commands.attribute_context.attribute_context_args import AttributeContextArgs | |
| from inseq.commands.attribute_context.attribute_context_helpers import ( | |
| AttributeContextOutput, | |
| filter_rank_tokens, | |
| get_filtered_tokens, | |
| ) | |
| from inseq.models import HuggingfaceModel | |
| def get_formatted_attribute_context_results( | |
| model: HuggingfaceModel, | |
| args: AttributeContextArgs, | |
| output: AttributeContextOutput, | |
| ) -> str: | |
| """Format the results of the context attribution process.""" | |
| def format_context_comment( | |
| model: HuggingfaceModel, | |
| has_other_context: bool, | |
| special_tokens_to_keep: list[str], | |
| context: str, | |
| context_scores: list[float], | |
| other_context_scores: Optional[list[float]] = None, | |
| is_target: bool = False, | |
| ) -> str: | |
| context_tokens = get_filtered_tokens( | |
| context, | |
| model, | |
| special_tokens_to_keep, | |
| replace_special_characters=True, | |
| is_target=is_target, | |
| ) | |
| context_token_tuples = [(t, None) for t in context_tokens] | |
| scores = context_scores | |
| if has_other_context: | |
| scores += other_context_scores | |
| context_ranked_tokens, _ = filter_rank_tokens( | |
| tokens=context_tokens, | |
| scores=scores, | |
| std_threshold=args.attribution_std_threshold, | |
| topk=args.attribution_topk, | |
| ) | |
| for idx, _, tok in context_ranked_tokens: | |
| context_token_tuples[idx] = (tok, "Influential context") | |
| return context_token_tuples | |
| out = [] | |
| output_current_tokens = get_filtered_tokens( | |
| output.output_current, | |
| model, | |
| args.special_tokens_to_keep, | |
| replace_special_characters=True, | |
| is_target=True, | |
| ) | |
| for example_idx, cci_out in enumerate(output.cci_scores, start=1): | |
| curr_output_tokens = [(t, None) for t in output_current_tokens] | |
| cti_idx = cci_out.cti_idx | |
| curr_output_tokens[cti_idx] = ( | |
| curr_output_tokens[cti_idx][0], | |
| "Context sensitive", | |
| ) | |
| if args.has_input_context: | |
| input_context_tokens = format_context_comment( | |
| model, | |
| args.has_output_context, | |
| args.special_tokens_to_keep, | |
| output.input_context, | |
| cci_out.input_context_scores, | |
| cci_out.output_context_scores, | |
| ) | |
| if args.has_output_context: | |
| output_context_tokens = format_context_comment( | |
| model, | |
| args.has_input_context, | |
| args.special_tokens_to_keep, | |
| output.output_context, | |
| cci_out.output_context_scores, | |
| cci_out.input_context_scores, | |
| is_target=True, | |
| ) | |
| out += [ | |
| ("\n\n" if example_idx > 1 else "", None), | |
| ( | |
| f"#{example_idx}.\nGenerated output:\t", | |
| None, | |
| ), | |
| ] | |
| out += curr_output_tokens | |
| if args.has_input_context: | |
| out += [("\nInput context:\t", None)] | |
| out += input_context_tokens | |
| if args.has_output_context: | |
| out += [("\nOutput context:\t", None)] | |
| out += output_context_tokens | |
| return out | |