diff --git a/captum/attr/__init__.py b/captum/attr/__init__.py index 19a3f5e9a5..93aa181a02 100644 --- a/captum/attr/__init__.py +++ b/captum/attr/__init__.py @@ -30,7 +30,11 @@ ) from captum.attr._core.layer.layer_lrp import LayerLRP # noqa from captum.attr._core.lime import Lime, LimeBase # noqa -from captum.attr._core.llm_attr import LLMAttribution, LLMGradientAttribution # noqa +from captum.attr._core.llm_attr import ( # noqa + LLMAttribution, + LLMAttributionResult, + LLMGradientAttribution, +) from captum.attr._core.lrp import LRP # noqa from captum.attr._core.neuron.neuron_conductance import NeuronConductance # noqa from captum.attr._core.neuron.neuron_deep_lift import ( # noqa @@ -67,6 +71,7 @@ NeuronAttribution, PerturbationAttribution, ) +from captum.attr._utils.baselines import ProductBaselines # noqa from captum.attr._utils.class_summarizer import ClassSummarizer from captum.attr._utils.interpretable_input import ( # noqa InterpretableInput, @@ -115,6 +120,7 @@ "LayerActivation", "LayerFeatureAblation", "LLMAttribution", + "LLMAttributionResult", "LLMGradientAttribution", "InternalInfluence", "InterpretableInput", @@ -134,6 +140,7 @@ "NeuronDeconvolution", "NeuronGuidedBackprop", "NoiseTunnel", + "ProductBaselines", "GradientShap", "InterpretableEmbeddingBase", "TextTemplateInput",