Skip to content

Commit 7dac7f4

Browse files
craymichaelfacebook-github-bot
authored andcommitted
Convert LLM Attribution result to a proper dataclass (#1572)
Summary: Pull Request resolved: #1572 Converts LLM attribution result to a proper dataclass. Reviewed By: aobo-y Differential Revision: D75727427 fbshipit-source-id: df983f47ed9c6a74516a1cf630ab8ef4276f0feb
1 parent 214acab commit 7dac7f4

File tree

1 file changed

+6
-14
lines changed

1 file changed

+6
-14
lines changed

captum/attr/_core/llm_attr.py

Lines changed: 6 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,9 @@
33
import warnings
44

55
from abc import ABC
6-
76
from copy import copy
8-
7+
from dataclasses import dataclass
98
from textwrap import shorten
10-
119
from typing import Any, Callable, cast, Dict, List, Optional, Tuple, Type, Union
1210

1311
import matplotlib.colors as mcolors
@@ -45,24 +43,18 @@
4543
}
4644

4745

46+
@dataclass
4847
class LLMAttributionResult:
4948
"""
5049
Data class for the return result of LLMAttribution,
5150
which includes the necessary properties of the attribution.
5251
It also provides utilities to help present and plot the result in different forms.
5352
"""
5453

55-
def __init__(
56-
self,
57-
seq_attr: Tensor,
58-
token_attr: Optional[Tensor],
59-
input_tokens: List[str],
60-
output_tokens: List[str],
61-
) -> None:
62-
self.seq_attr = seq_attr
63-
self.token_attr = token_attr
64-
self.input_tokens = input_tokens
65-
self.output_tokens = output_tokens
54+
seq_attr: Tensor
55+
token_attr: Optional[Tensor]
56+
input_tokens: List[str]
57+
output_tokens: List[str]
6658

6759
@property
6860
def seq_attr_dict(self) -> Dict[str, float]:

0 commit comments

Comments
 (0)