13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191 | class AimCallback(BaseCallbackHandler):
"""
AimCallback callback class.
Args:
repo (:obj:`str`, optional):
Aim repository path or Repo object to which Run object is bound.
If skipped, default Repo is used.
experiment_name (:obj:`str`, optional):
Sets Run's `experiment` property. 'default' if not specified.
Can be used later to query runs/sequences.
system_tracking_interval (:obj:`int`, optional):
Sets the tracking interval in seconds for system usage
metrics (CPU, Memory, etc.). Set to `None` to disable
system metrics tracking.
log_system_params (:obj:`bool`, optional):
Enable/Disable logging of system params such as installed packages,
git info, environment variables, etc.
capture_terminal_logs (:obj:`bool`, optional):
Enable/Disable terminal stdout logging.
event_starts_to_ignore (Optional[List[CBEventType]]):
list of event types to ignore when tracking event starts.
event_ends_to_ignore (Optional[List[CBEventType]]):
list of event types to ignore when tracking event ends.
"""
def __init__(
self,
repo: Optional[str] = None,
experiment_name: Optional[str] = None,
system_tracking_interval: Optional[int] = 1,
log_system_params: Optional[bool] = True,
capture_terminal_logs: Optional[bool] = True,
event_starts_to_ignore: Optional[List[CBEventType]] = None,
event_ends_to_ignore: Optional[List[CBEventType]] = None,
run_params: Optional[Dict[str, Any]] = None,
) -> None:
if Run is None:
raise ModuleNotFoundError(
"Please install aim to use the AimCallback: 'pip install aim'"
)
event_starts_to_ignore = (
event_starts_to_ignore if event_starts_to_ignore else []
)
event_ends_to_ignore = event_ends_to_ignore if event_ends_to_ignore else []
super().__init__(
event_starts_to_ignore=event_starts_to_ignore,
event_ends_to_ignore=event_ends_to_ignore,
)
self.repo = repo
self.experiment_name = experiment_name
self.system_tracking_interval = system_tracking_interval
self.log_system_params = log_system_params
self.capture_terminal_logs = capture_terminal_logs
self._run: Optional[Any] = None
self._run_hash = None
self._llm_response_step = 0
self.setup(run_params)
def on_event_start(
self,
event_type: CBEventType,
payload: Optional[Dict[str, Any]] = None,
event_id: str = "",
parent_id: str = "",
**kwargs: Any,
) -> str:
"""
Args:
event_type (CBEventType): event type to store.
payload (Optional[Dict[str, Any]]): payload to store.
event_id (str): event id to store.
parent_id (str): parent event id.
"""
return ""
def on_event_end(
self,
event_type: CBEventType,
payload: Optional[Dict[str, Any]] = None,
event_id: str = "",
**kwargs: Any,
) -> None:
"""
Args:
event_type (CBEventType): event type to store.
payload (Optional[Dict[str, Any]]): payload to store.
event_id (str): event id to store.
"""
if not self._run:
raise ValueError("AimCallback failed to init properly.")
if event_type is CBEventType.LLM and payload:
if EventPayload.PROMPT in payload:
llm_input = str(payload[EventPayload.PROMPT])
llm_output = str(payload[EventPayload.COMPLETION])
else:
message = payload.get(EventPayload.MESSAGES, [])
llm_input = "\n".join([str(x) for x in message])
llm_output = str(payload[EventPayload.RESPONSE])
self._run.track(
Text(llm_input),
name="prompt",
step=self._llm_response_step,
context={"event_id": event_id},
)
self._run.track(
Text(llm_output),
name="response",
step=self._llm_response_step,
context={"event_id": event_id},
)
self._llm_response_step += 1
elif event_type is CBEventType.CHUNKING and payload:
for chunk_id, chunk in enumerate(payload[EventPayload.CHUNKS]):
self._run.track(
Text(chunk),
name="chunk",
step=self._llm_response_step,
context={"chunk_id": chunk_id, "event_id": event_id},
)
@property
def experiment(self) -> Run:
if not self._run:
self.setup()
return self._run
def setup(self, args: Optional[Dict[str, Any]] = None) -> None:
if not self._run:
if self._run_hash:
self._run = Run(
self._run_hash,
repo=self.repo,
system_tracking_interval=self.system_tracking_interval,
log_system_params=self.log_system_params,
capture_terminal_logs=self.capture_terminal_logs,
)
else:
self._run = Run(
repo=self.repo,
experiment=self.experiment_name,
system_tracking_interval=self.system_tracking_interval,
log_system_params=self.log_system_params,
capture_terminal_logs=self.capture_terminal_logs,
)
self._run_hash = self._run.hash
# Log config parameters
if args:
try:
for key in args:
self._run.set(key, args[key], strict=False)
except Exception as e:
logger.warning(f"Aim could not log config parameters -> {e}")
def __del__(self) -> None:
if self._run and self._run.active:
self._run.close()
def start_trace(self, trace_id: Optional[str] = None) -> None:
pass
def end_trace(
self,
trace_id: Optional[str] = None,
trace_map: Optional[Dict[str, List[str]]] = None,
) -> None:
pass
|
