Add treescope requirement

inseq-team · Aug 1, 2024 · e03ba89 · e03ba89
1 parent e5b835b
commit e03ba89
Show file tree

Hide file tree

Showing 5 changed files with 11 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,8 @@
 
 ## 🚀 Features
 
+- Added [treescope](https://github.com/google-deepmind/treescope) for model and tensor visualization.
+
 - Added new models `DbrxForCausalLM`, `OlmoForCausalLM`, `Phi3ForCausalLM`, `Qwen2MoeForCausalLM`, `Gemma2ForCausalLM` to model config.
 
 - Add `rescale_attributions` to Inseq CLI commands for `rescale=True` ([#280](https://github.com/inseq-team/inseq/pull/280)).

diff --git a/inseq/data/attribution.py b/inseq/data/attribution.py
@@ -290,8 +290,8 @@ def from_step_attributions(
             pos_start.append(curr_pos_start)
             source = tokenized_target_sentences[seq_idx][:curr_pos_start] if not sources else sources[seq_idx]
             curr_seq_attribution: FeatureAttributionSequenceOutput = attr.get_sequence_cls(
-                source=source,
-                target=tokenized_target_sentences[seq_idx],
+                source=deepcopy(source),
+                target=deepcopy(tokenized_target_sentences[seq_idx]),
                 attr_pos_start=pos_start[seq_idx],
                 attr_pos_end=attr_pos_end,
             )

diff --git a/pyproject.toml b/pyproject.toml
@@ -50,6 +50,7 @@ dependencies = [
   "torch>=2.0",
   "matplotlib>=3.5.3",
   "tqdm>=4.64.0",
+  "treescope>=0.1.0",
   "nvidia-cublas-cu11>=11.10.3.66; sys_platform=='Linux'",
   "nvidia-cuda-cupti-cu11>=11.7.101; sys_platform=='Linux'",
   "nvidia-cuda-nvrtc-cu11>=11.7.99; sys_platform=='Linux'",

diff --git a/requirements-dev.txt b/requirements-dev.txt
@@ -190,6 +190,7 @@ numpy==1.26.4
     #   scikit-learn
     #   scipy
     #   transformers
+    #   treescope
 packaging==23.2
     # via
     #   datasets
@@ -402,6 +403,8 @@ traitlets==5.14.1
     #   matplotlib-inline
 transformers==4.38.1
     # via inseq (pyproject.toml)
+treescope==0.1.0
+    # via inseq (pyproject.toml)
 typeguard==2.13.3
     # via
     #   inseq (pyproject.toml)

diff --git a/requirements.txt b/requirements.txt
@@ -55,6 +55,7 @@ numpy==1.26.4
     #   jaxtyping
     #   matplotlib
     #   transformers
+    #   treescope
 packaging==23.2
     # via
     #   huggingface-hub
@@ -106,6 +107,8 @@ tqdm==4.66.4
     #   transformers
 transformers==4.38.1
     # via inseq (pyproject.toml)
+treescope==0.1.0
+    # via inseq (pyproject.toml)
 typeguard==2.13.3
     # via
     #   inseq (pyproject.toml)