24 add cross attention labels text #60

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

meilame-tayebjee merged 22 commits into main from 24-add-cross-attention-labels-text

Jan 27, 2026

+408 −75

.gitignore

-Original file line number
+Diff line change
@@ Expand Up / @@ -183,3 +183,4 @@ example_files/ @@
     _site/
     .quarto/
     **/*.quarto_ipynb
+    my_ttc/

tests/test_pipeline.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -9,6 +9,7 @@ @@
         AttentionConfig,
         CategoricalVariableNet,
         ClassificationHead,
+        LabelAttentionConfig,
         TextEmbedder,
         TextEmbedderConfig,
     )
@@ Expand Down Expand Up / @@ -51,7 +52,14 @@ def model_params(): @@
         }
-    def run_full_pipeline(tokenizer, sample_text_data, categorical_data, labels, model_params):
+    def run_full_pipeline(
+        tokenizer,
+        sample_text_data,
+        categorical_data,
+        labels,
+        model_params,
+        label_attention_enabled: bool = False,
+    ):
         """Helper function to run the complete pipeline for a given tokenizer."""
         # Create dataset
         dataset = TextClassificationDataset(
@@ Expand Down Expand Up @@
             embedding_dim=model_params["embedding_dim"],
             padding_idx=padding_idx,
             attention_config=attention_config,
+            label_attention_config=(
+                LabelAttentionConfig(
+                    n_head=attention_config.n_head,
+                    num_classes=model_params["num_classes"],
+                )
+                if label_attention_enabled
+                else None
+            ),
         )
         text_embedder = TextEmbedder(text_embedder_config=text_embedder_config)
@@ Expand All @@
         expected_input_dim = model_params["embedding_dim"] + categorical_var_net.output_dim
         classification_head = ClassificationHead(
             input_dim=expected_input_dim,
-            num_classes=model_params["num_classes"],
+            num_classes=model_params["num_classes"] if not label_attention_enabled else 1,
         )
         # Create model
@@ Expand Down Expand Up @@
             categorical_embedding_dims=model_params["categorical_embedding_dims"],
             num_classes=model_params["num_classes"],
             attention_config=attention_config,
+            label_attention_config=(
+                LabelAttentionConfig(
+                    n_head=attention_config.n_head,
+                    num_classes=model_params["num_classes"],
+                )
+                if label_attention_enabled
+                else None
+            ),
         )
         # Create training config
@@ Expand Down Expand Up @@
         # Predict with explanations
         top_k = 5
-        predictions = ttc.predict(X, top_k=top_k, explain=True)
+        predictions = ttc.predict(
+            X,
+            top_k=top_k,
+            explain_with_label_attention=label_attention_enabled,
+            explain_with_captum=True,
+        )
+        # Test label attention assertions
+        if label_attention_enabled:
+            assert (
+                predictions["label_attention_attributions"] is not None
+            ), "Label attention attributions should not be None when label_attention_enabled is True"
+            label_attention_attributions = predictions["label_attention_attributions"]
+            expected_shape = (
+                len(sample_text_data),  # batch_size
+                model_params["n_head"],  # n_head
+                model_params["num_classes"],  # num_classes
+                tokenizer.output_dim,  # seq_len
+            )
+            assert label_attention_attributions.shape == expected_shape, (
+                f"Label attention attributions shape mismatch. "
+                f"Expected {expected_shape}, got {label_attention_attributions.shape}"
+            )
+        else:
+            # When label attention is not enabled, the attributions should be None
+            assert (
+                predictions.get("label_attention_attributions") is None
+            ), "Label attention attributions should be None when label_attention_enabled is False"
         # Test explainability functions
         text_idx = 0
         text = sample_text_data[text_idx]
         offsets = predictions["offset_mapping"][text_idx]
-        attributions = predictions["attributions"][text_idx]
+        attributions = predictions["captum_attributions"][text_idx]
         word_ids = predictions["word_ids"][text_idx]
         words, word_attributions = map_attributions_to_word(attributions, text, word_ids, offsets)
@@ Expand Down Expand Up / @@ -239,3 +291,26 @@ def test_ngram_tokenizer(sample_data, model_params): @@
         # Run full pipeline
         run_full_pipeline(tokenizer, sample_text_data, categorical_data, labels, model_params)
+    def test_label_attention_enabled(sample_data, model_params):
+        """Test the full pipeline with label attention enabled (using WordPieceTokenizer)."""
+        sample_text_data, categorical_data, labels = sample_data
+        vocab_size = 100
+        tokenizer = WordPieceTokenizer(vocab_size, output_dim=50)
+        tokenizer.train(sample_text_data)
+        # Check tokenizer works
+        result = tokenizer.tokenize(sample_text_data)
+        assert result.input_ids.shape[0] == len(sample_text_data)
+        # Run full pipeline with label attention enabled
+        run_full_pipeline(
+            tokenizer,
+            sample_text_data,
+            categorical_data,
+            labels,
+            model_params,
+            label_attention_enabled=True,
+        )

torchTextClassifiers/model/components/__init__.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -8,5 +8,6 @@ @@
         CategoricalVariableNet as CategoricalVariableNet,
     )
     from .classification_head import ClassificationHead as ClassificationHead
+    from .text_embedder import LabelAttentionConfig as LabelAttentionConfig
     from .text_embedder import TextEmbedder as TextEmbedder
     from .text_embedder import TextEmbedderConfig as TextEmbedderConfig

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

24 add cross attention labels text #60

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!