Fix attribute-context for current preceding context in input

inseq-team · Feb 27, 2024 · 8e85643 · 8e85643
1 parent b038877
commit 8e85643
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/inseq/commands/attribute_context/attribute_context.py b/inseq/commands/attribute_context/attribute_context.py
@@ -206,6 +206,7 @@ def attribute_context_with_model(args: AttributeContextArgs, model: HuggingfaceM
             model,
             cci_attrib_out,
             args.input_template,
+            args.input_current_text,
             input_context_tokens,
             input_full_tokens,
             args.output_template,

diff --git a/inseq/commands/attribute_context/attribute_context_helpers.py b/inseq/commands/attribute_context/attribute_context_helpers.py
@@ -402,6 +402,7 @@ def get_source_target_cci_scores(
     model: HuggingfaceModel,
     cci_attrib_out: FeatureAttributionSequenceOutput,
     input_template: str,
+    input_current_text: str,
     input_context_tokens: list[str],
     input_full_tokens: list[str],
     output_template: str,
@@ -421,6 +422,8 @@ def get_source_target_cci_scores(
         else:
             input_scores = cci_attrib_out.target_attributions[:, 0].tolist()
         input_prefix, *_ = input_template.partition("{context}")
+        if "{current}" in input_prefix:
+            input_prefix = input_prefix.format(current=input_current_text)
         input_prefix_tokens = get_filtered_tokens(input_prefix, model, special_tokens_to_keep, is_target=False)
         input_prefix_len = len(input_prefix_tokens)
         input_scores = input_scores[input_prefix_len : len(input_context_tokens) + input_prefix_len]