From b1d8dac8c246f19547ae0cc044076d40633fab2a Mon Sep 17 00:00:00 2001
From: Gabriele Sarti <gabriele.sarti996@gmail.com>
Date: Tue, 13 Aug 2024 18:24:51 +0200
Subject: [PATCH] Fix tests

---
 README.md                                | 2 +-
 tests/commands/test_attribute_context.py | 7 +++++++
 2 files changed, 8 insertions(+), 1 deletion(-)

diff --git a/README.md b/README.md
index bb5af86..50f066e 100644
--- a/README.md
+++ b/README.md
@@ -252,7 +252,7 @@ All commands support the full range of parameters available for `attribute`, att
 
   **Result:**
 
-  <img src="https://raw.githubusercontent.com/inseq-team/inseq/main/docs/source/images/attribute_context_hospital_output.png" style="width:300px">
+  <img src="https://raw.githubusercontent.com/inseq-team/inseq/main/docs/source/images/attribute_context_hospital_output.png" style="width:500px">
 </details>
 
 ## Planned Development
diff --git a/tests/commands/test_attribute_context.py b/tests/commands/test_attribute_context.py
index 1013b1c..a0aa05c 100644
--- a/tests/commands/test_attribute_context.py
+++ b/tests/commands/test_attribute_context.py
@@ -59,6 +59,7 @@ def test_in_out_ctx_encdec_whitespace_sep(encdec_model: MarianMTModel):
             CCIOutput(
                 cti_idx=0,
                 cti_token="▁Où",
+                contrast_token="▁Où",
                 cti_score=1.36,
                 contextual_output="Où",
                 contextless_output="Où",
@@ -95,6 +96,7 @@ def test_in_ctx_deconly(deconly_model: GPT2LMHeadModel):
             CCIOutput(
                 cti_idx=2,
                 cti_token="Ġhospital",
+                contrast_token="Ġoffice",
                 cti_score=0.55,
                 contextual_output="George was sick yesterday. His colleagues asked him to come to the hospital",
                 contextless_output="His colleagues asked him to come to the office",
@@ -164,6 +166,7 @@ def test_out_ctx_deconly(deconly_model: GPT2LMHeadModel):
             CCIOutput(
                 cti_idx=1,
                 cti_token="Ġpairs",
+                contrast_token="Ġhorses",
                 cti_score=1.39,
                 contextual_output="Question: How many pairs of legs do 10 horses have?\nLet's think step by step:\n1. A horse has 4 legs.\n2. 10 horses have 40 legs.\n3. 40 legs make 20 pairs of legs.\n\nAnswer:\n20 pairs",
                 contextless_output="Question: How many pairs of legs do 10 horses have?\nAnswer:\n20 horses",
@@ -230,6 +233,7 @@ def test_in_out_ctx_deconly(deconly_model: GPT2LMHeadModel):
             CCIOutput(
                 cti_idx=2,
                 cti_token="Ġfine",
+                contrast_token="Ġa",
                 cti_score=1.5,
                 contextual_output="George was sick yesterday. His colleagues asked him if something was wrong. He said he was fine",
                 contextless_output="His colleagues asked him if he was a",
@@ -269,6 +273,7 @@ def test_in_ctx_encdec_special_sep():
             CCIOutput(
                 cti_idx=3,
                 cti_token="elles",
+                contrast_token="elles",
                 cti_score=0.32,
                 contextual_output="Où sont-elles",
                 contextless_output="Où sont-elles",
@@ -310,6 +315,7 @@ def test_in_out_ctx_encdec_special_sep():
             CCIOutput(
                 cti_idx=3,
                 cti_token="elles",
+                contrast_token="ils",
                 cti_score=3.99,
                 contextual_output="Les filles étaient parties.<brk>  Où sont-elles",
                 contextless_output="Où sont-ils",
@@ -352,6 +358,7 @@ def test_in_out_ctx_encdec_langtag_whitespace_sep():
             CCIOutput(
                 cti_idx=4,
                 cti_token="elles",
+                contrast_token="ils",
                 cti_score=4.49,
                 contextual_output="Les filles étaient loin. Où sont-elles",
                 contextless_output="Où sont-ils",