ICV Aligner Bugs

[abhishekreddyallu]

ICV workflows have device, adapter-layout, and threshold-forwarding bugs

1. CPU configuration unconditionally transfers ICVs to CUDA

Location: ontoaligner/aligner/icv/icv.py:225

Current code:

 torch.stack(icvs_to_shift, dim=1).cuda(), alpha=[self.icv_alpha]

This forces ICV tensors onto CUDA even when the user configures CPU execution.

Observed before fix:

AssertionError: Torch not compiled with CUDA enabled

This happens in CPU-only Torch environments because .cuda() is called unconditionally.

Expected behavior:

ICV tensors should remain on the configured device.

Fix:

torch.stack(icvs_to_shift, dim=1).to(self.kwargs.get("device", "cpu")), alpha=[self.icv_alpha]

Verified after fix:

ICV tensor device: cuda
PASS: build_icv respects device='cuda'

ICV tensor device: cpu
PASS: build_icv respects device='cpu'

---

2. ICVAdapter assumes Falcon decoder layout

Location: ontoaligner/aligner/icv/icv.py:158-174

Current code:

            self.model.transformer.h[i].mlp = torch.nn.Sequential(
                self.model.transformer.h[i].mlp, AdapterLayer(icvs_, alpha)
            )
        return self.model

Adapter removal also assumes:

self.model.transformer.h[i].mlp

This works for Falcon-style models, but fails for other decoder layouts.

Observed before fix:

Falcon-like layout
PASS: adapter applied
PASS: adapter removed

LLaMA/Vicuna-like layout
FAIL: AttributeError: 'LlamaLikeModel' object has no attribute 'transformer'

MPT-like layout
FAIL: AttributeError: 'MPTTransformer' object has no attribute 'h'

Root cause:

The adapter is hard-coded to Falcon’s decoder structure:

model.transformer.h[i].mlp

But other supported model families use different layouts:

LLaMA / Vicuna -> model.model.layers
MPT            -> model.transformer.blocks
Falcon         -> model.transformer.h

Fix:

Use the existing architecture-aware layer helper and dynamically replace each layer’s MLP/FFN module.

from ..llm.utils.llm_layers import get_layers

    def _get_mlp_attr_name(self, layer):
        """
        Returns the feed-forward/MLP attribute name used by the transformer layer.
        """
        for attr_name in ("mlp", "ffn", "feed_forward", "feedforward"):
            if hasattr(layer, attr_name):
                return attr_name

        raise AttributeError(
            f"Could not find MLP/FFN module in layer {layer.__class__.__name__}"
        )

    def get_model(self, icvs, alpha):
        """
        Adds ICV-based adapter layers to the model for ICV-based embedding adjustment.

        Parameters:
            icvs : list of torch.Tensor
                List of ICVs used for embedding adjustment.
            alpha : list of float
                List of scaling factors for ICV influence.

        Returns:
            torch.nn.Module
                The model with ICV-based adapter layers integrated.
        """
        layers = get_layers(self.model)

        for i, layer in enumerate(layers):
            mlp_attr_name = self._get_mlp_attr_name(layer)
            mlp_layer = getattr(layer, mlp_attr_name)

            setattr(layer,mlp_attr_name,torch.nn.Sequential(mlp_layer, AdapterLayer(icvs[i], alpha)))
        return self.model

    def remove_adapter(self):
        """
        Removes adapter layers from the model and restores the original architecture.
        """
        weight_all = []
        layers = get_layers(self.model)

        for layer in layers:
            mlp_attr_name = self._get_mlp_attr_name(layer)
            mlp_layer = getattr(layer, mlp_attr_name)

            weight_all.append(mlp_layer[1].weight_all)
            setattr(layer, mlp_attr_name, mlp_layer[0])

---

3. ICV generation ignores the configured retrieval threshold

Location: ontoaligner/aligner/icv/icv.py:281

Current code:

        ir_output_cleaned = process.retriever_postprocessor(predicts=ir_output)

This does not forward the configured retrieval threshold.

The base RAG implementation already forwards the threshold, but the ICV override does not.

Before fix:

LLM output: [{'candidate_count_seen_by_llm': 2, 'ir_output_seen_by_llm': [{'source': 'SourceConcept', 'target-cands': ['TargetA', 'TargetB'], 'score-cands': [0.95, 0.8]}]}]
Candidate count seen by LLM: 2
AssertionError: BUG: threshold=0.99 was not forwarded. Candidates with scores 0.95 and 0.80 still reached the LLM.

Fix:

        if 'threshold' in self.kwargs['retriever_config']:
            threshold = self.kwargs['retriever_config']['threshold']
        else:
            threshold = 0.0
        ir_output_cleaned = process.retriever_postprocessor(predicts=ir_output, threshold=threshold)
        

Verified after fix:

LLM output: [{'candidate_count_seen_by_llm': 0, 'ir_output_seen_by_llm': [{'source': 'SourceConcept', 'target-cands': [], 'score-cands': []}]}]
Candidate count seen by LLM: 0
PASS: ICV.generate() forwarded the configured retrieval threshold.