psaiops 0.3.2__py3-none-any.whl → 0.4.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- psaiops/combine/app.py +366 -0
- psaiops/common/data.py +8 -0
- psaiops/common/model.py +45 -0
- psaiops/common/tokenizer.py +41 -0
- psaiops/compose/contrast/app.py +4 -2
- psaiops/compose/contrast/lib.py +3 -38
- psaiops/compose/maths/app.py +8 -5
- psaiops/compose/maths/lib.py +0 -41
- psaiops/score/attention/app.py +7 -5
- psaiops/score/attention/lib.py +7 -82
- psaiops/score/residual/app.py +290 -0
- psaiops/score/residual/lib.py +134 -0
- psaiops/score/router/app.py +161 -13
- psaiops/score/router/lib.py +50 -57
- psaiops/score/shapley/app.py +5 -4
- psaiops/score/shapley/lib.py +0 -65
- psaiops/score/similarity/__init__.py +0 -0
- psaiops/score/similarity/app.py +152 -0
- psaiops/score/similarity/lib.py +1 -0
- {psaiops-0.3.2.dist-info → psaiops-0.4.0.dist-info}/METADATA +14 -19
- psaiops-0.4.0.dist-info/RECORD +36 -0
- {psaiops-0.3.2.dist-info → psaiops-0.4.0.dist-info}/WHEEL +1 -1
- psaiops-0.4.0.dist-info/licenses/.github/LICENSE.md +661 -0
- psaiops/common/dropdown.py +0 -19
- psaiops-0.3.2.dist-info/RECORD +0 -28
- /psaiops/{steer → score/residual}/__init__.py +0 -0
psaiops/score/router/app.py
CHANGED
|
@@ -3,13 +3,16 @@ import functools
|
|
|
3
3
|
import gradio
|
|
4
4
|
import torch
|
|
5
5
|
import torch.cuda
|
|
6
|
+
import matplotlib.pyplot
|
|
6
7
|
|
|
8
|
+
import psaiops.common.model
|
|
9
|
+
import psaiops.common.tokenizer
|
|
7
10
|
import psaiops.score.router.lib
|
|
8
11
|
|
|
9
12
|
# META #########################################################################
|
|
10
13
|
|
|
11
14
|
STYLE = '''.white-text span { color: white; }'''
|
|
12
|
-
TITLE = '''
|
|
15
|
+
TITLE = '''Router Scoring'''
|
|
13
16
|
INTRO = '''Plot the logits of the router for a given prompt.\nUnder construction, only "openai/gpt-oss-20b" is available for now.'''
|
|
14
17
|
|
|
15
18
|
MODEL = 'openai/gpt-oss-20b'
|
|
@@ -18,8 +21,8 @@ MODEL = 'openai/gpt-oss-20b'
|
|
|
18
21
|
|
|
19
22
|
def create_color_map() -> dict:
|
|
20
23
|
return {
|
|
21
|
-
'
|
|
22
|
-
|
|
24
|
+
'0': '#000000',
|
|
25
|
+
'1': '#004444',}
|
|
23
26
|
|
|
24
27
|
# INTRO ########################################################################
|
|
25
28
|
|
|
@@ -50,16 +53,23 @@ def create_inputs_block() -> dict:
|
|
|
50
53
|
__input = gradio.Textbox(label='Prompt', value='', placeholder='A string of tokens to score.', lines=4, scale=1, show_copy_button=True, interactive=True)
|
|
51
54
|
return {'input_block': __input}
|
|
52
55
|
|
|
56
|
+
# PLOTS ########################################################################
|
|
57
|
+
|
|
58
|
+
def create_plot_block() -> dict:
|
|
59
|
+
__plot = gradio.Plot(label='Router', scale=1)
|
|
60
|
+
return {'plot_block': __plot,}
|
|
61
|
+
|
|
53
62
|
# OUTPUTS ######################################################################
|
|
54
63
|
|
|
55
64
|
def create_outputs_block() -> dict:
|
|
56
|
-
__output = gradio.
|
|
57
|
-
return {'output_block': __output
|
|
65
|
+
__output = gradio.HighlightedText(label='Output', value='', scale=1, interactive=False, show_legend=False, show_inline_category=False, combine_adjacent=False, color_map=create_color_map(), elem_classes='white-text')
|
|
66
|
+
return {'output_block': __output}
|
|
58
67
|
|
|
59
68
|
# SELECT #######################################################################
|
|
60
69
|
|
|
61
70
|
def create_selection_block() -> dict:
|
|
62
|
-
|
|
71
|
+
# __play = gradio.Button('>', variant='primary', size='lg', scale=1, interactive=True)
|
|
72
|
+
__position = gradio.Slider(label='Token', value=-1, minimum=-1, maximum=15, step=1, scale=1, interactive=True) # info='-1 to average on all tokens'
|
|
63
73
|
return {'position_block': __position,}
|
|
64
74
|
|
|
65
75
|
# ACTIONS ######################################################################
|
|
@@ -71,7 +81,9 @@ def create_actions_block() -> dict:
|
|
|
71
81
|
# STATE ########################################################################
|
|
72
82
|
|
|
73
83
|
def create_state() -> dict:
|
|
74
|
-
return {
|
|
84
|
+
return {
|
|
85
|
+
'output_state': gradio.State(None),
|
|
86
|
+
'router_state': gradio.State(None),}
|
|
75
87
|
|
|
76
88
|
# LAYOUT #######################################################################
|
|
77
89
|
|
|
@@ -83,6 +95,8 @@ def create_layout(intro: str=INTRO) -> dict:
|
|
|
83
95
|
__fields.update({'main_tab': __main_tab})
|
|
84
96
|
with gradio.Row(equal_height=True):
|
|
85
97
|
__fields.update(create_inputs_block())
|
|
98
|
+
with gradio.Row(equal_height=True):
|
|
99
|
+
__fields.update(create_plot_block())
|
|
86
100
|
with gradio.Row(equal_height=True):
|
|
87
101
|
__fields.update(create_outputs_block())
|
|
88
102
|
with gradio.Row(equal_height=True):
|
|
@@ -100,8 +114,100 @@ def create_layout(intro: str=INTRO) -> dict:
|
|
|
100
114
|
|
|
101
115
|
# EVENTS #######################################################################
|
|
102
116
|
|
|
103
|
-
def update_position_range(
|
|
104
|
-
|
|
117
|
+
def update_position_range(
|
|
118
|
+
current_val: float,
|
|
119
|
+
token_num: float,
|
|
120
|
+
output_data: torch.Tensor,
|
|
121
|
+
) -> dict:
|
|
122
|
+
# take the generated tokens into account
|
|
123
|
+
__max = int(token_num) - 1 if (output_data is None) else int(output_data.shape[-1])
|
|
124
|
+
# keep the previous value if possible
|
|
125
|
+
__val = min(int(current_val), __max)
|
|
126
|
+
# return a gradio update dictionary
|
|
127
|
+
return gradio.update(maximum=__max, value=__val)
|
|
128
|
+
|
|
129
|
+
def update_computation_state(
|
|
130
|
+
token_num: float,
|
|
131
|
+
topk_num: float,
|
|
132
|
+
topp_num: float,
|
|
133
|
+
token_idx: float,
|
|
134
|
+
prompt_str: str,
|
|
135
|
+
device_str: str,
|
|
136
|
+
model_obj: object,
|
|
137
|
+
tokenizer_obj: object,
|
|
138
|
+
) -> tuple:
|
|
139
|
+
# sanitize the inputs
|
|
140
|
+
__token_num = max(1, min(128, int(token_num)))
|
|
141
|
+
__topk_num = max(1, min(8, int(topk_num)))
|
|
142
|
+
__topp_num = max(0.0, min(1.0, float(topp_num)))
|
|
143
|
+
__token_idx = max(-1, min(__token_num, int(token_idx)))
|
|
144
|
+
__prompt_str = prompt_str.strip()
|
|
145
|
+
__device_str = device_str if (device_str in ['cpu', 'cuda']) else 'cpu'
|
|
146
|
+
# exit if some values are missing
|
|
147
|
+
if (not __prompt_str) or (model_obj is None) or (tokenizer_obj is None):
|
|
148
|
+
return (torch.empty(0), torch.empty(0))
|
|
149
|
+
# dictionary {'input_ids': _, 'attention_mask': _}
|
|
150
|
+
__input_data = psaiops.common.tokenizer.preprocess_token_ids(
|
|
151
|
+
tokenizer_obj=tokenizer_obj,
|
|
152
|
+
prompt_str=__prompt_str,
|
|
153
|
+
device_str=__device_str)
|
|
154
|
+
# tensor (1, T)
|
|
155
|
+
__output_data = psaiops.common.model.generate_token_ids(
|
|
156
|
+
model_obj=model_obj,
|
|
157
|
+
input_args=__input_data,
|
|
158
|
+
token_num=__token_num,
|
|
159
|
+
topk_num=__topk_num,
|
|
160
|
+
topp_num=__topp_num)
|
|
161
|
+
# tensor (L, S, H, T, T)
|
|
162
|
+
__router_data = psaiops.score.router.lib.compute_router_weights(
|
|
163
|
+
model_obj=model_obj,
|
|
164
|
+
token_data=__output_data)
|
|
165
|
+
# update each component => (highlight, plot) states
|
|
166
|
+
return (
|
|
167
|
+
__output_data.cpu(),
|
|
168
|
+
__router_data.cpu(),)
|
|
169
|
+
|
|
170
|
+
def update_router_plot(
|
|
171
|
+
token_idx: float,
|
|
172
|
+
router_data: torch.Tensor,
|
|
173
|
+
) -> tuple:
|
|
174
|
+
# exit if some values are missing
|
|
175
|
+
if (router_data is None) or (len(router_data) == 0):
|
|
176
|
+
return None
|
|
177
|
+
# reduce the batch and token axes => tensor (L, E)
|
|
178
|
+
__plot_data = psaiops.score.router.lib.reduce_router_weights(
|
|
179
|
+
router_data=router_data,
|
|
180
|
+
token_idx=int(token_idx),)
|
|
181
|
+
# translate the scores into integer labels
|
|
182
|
+
__plot_data = psaiops.score.router.lib.postprocess_router_weights(
|
|
183
|
+
router_data=__plot_data,)
|
|
184
|
+
# plot the data
|
|
185
|
+
__figure, __axes = matplotlib.pyplot.subplots()
|
|
186
|
+
__axes.imshow(__plot_data.float().numpy(), vmin=0.0, vmax=1.0, cmap='viridis')
|
|
187
|
+
__figure.tight_layout()
|
|
188
|
+
# remove the figure for the pyplot register for garbage collection
|
|
189
|
+
matplotlib.pyplot.close(__figure)
|
|
190
|
+
# update each component => (highlight, plot) states
|
|
191
|
+
return __figure
|
|
192
|
+
|
|
193
|
+
def update_text_highlight(
|
|
194
|
+
token_idx: float,
|
|
195
|
+
output_data: torch.Tensor,
|
|
196
|
+
tokenizer_obj: object,
|
|
197
|
+
) -> list:
|
|
198
|
+
# exit if some values are missing
|
|
199
|
+
if (output_data is None) or (len(output_data) == 0):
|
|
200
|
+
return None
|
|
201
|
+
# detokenize the IDs
|
|
202
|
+
__token_str = psaiops.common.tokenizer.postprocess_token_ids(
|
|
203
|
+
tokenizer_obj=tokenizer_obj,
|
|
204
|
+
token_data=output_data)
|
|
205
|
+
# list of string classes
|
|
206
|
+
__token_cls = psaiops.score.router.lib.postprocess_token_cls(
|
|
207
|
+
token_idx=int(token_idx),
|
|
208
|
+
token_dim=len(__token_str))
|
|
209
|
+
# pairs of token and class
|
|
210
|
+
return list(zip(__token_str, __token_cls))
|
|
105
211
|
|
|
106
212
|
# APP ##########################################################################
|
|
107
213
|
|
|
@@ -110,19 +216,61 @@ def create_app(title: str=TITLE, intro: str=INTRO, style: str=STYLE, model: str=
|
|
|
110
216
|
with gradio.Blocks(theme=gradio.themes.Soft(), title=title, css=style) as __app:
|
|
111
217
|
# load the model
|
|
112
218
|
__device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
113
|
-
|
|
114
|
-
__tokenizer = psaiops.
|
|
219
|
+
__model = psaiops.common.model.get_model(name=model, device=__device)
|
|
220
|
+
__tokenizer = psaiops.common.tokenizer.get_tokenizer(name=model, device=__device)
|
|
221
|
+
# adapt the event handlers
|
|
222
|
+
__compute = functools.partial(update_computation_state, model_obj=__model, tokenizer_obj=__tokenizer, device_str=__device)
|
|
223
|
+
__highlight = functools.partial(update_text_highlight, tokenizer_obj=__tokenizer)
|
|
115
224
|
# create the UI
|
|
116
225
|
__fields.update(create_layout(intro=intro))
|
|
117
226
|
# init the state
|
|
118
227
|
__fields.update(create_state())
|
|
119
|
-
#
|
|
228
|
+
# update the data after clicking process
|
|
229
|
+
__fields['process_block'].click(
|
|
230
|
+
fn=__compute,
|
|
231
|
+
inputs=[__fields[__k] for __k in ['tokens_block', 'topk_block', 'topp_block', 'position_block', 'input_block']],
|
|
232
|
+
outputs=[__fields[__k] for __k in ['output_state', 'router_state']],
|
|
233
|
+
queue=False,
|
|
234
|
+
show_progress='full').then(
|
|
235
|
+
# update the range of the position slider when the output changes
|
|
236
|
+
fn=update_position_range,
|
|
237
|
+
inputs=[__fields[__k] for __k in ['position_block', 'tokens_block', 'output_state']],
|
|
238
|
+
outputs=__fields['position_block'],
|
|
239
|
+
queue=False,
|
|
240
|
+
show_progress='hidden').then(
|
|
241
|
+
# update the token highlight when the output data changes
|
|
242
|
+
fn=__highlight,
|
|
243
|
+
inputs=[__fields[__k] for __k in ['position_block', 'output_state']],
|
|
244
|
+
outputs=__fields['output_block'],
|
|
245
|
+
queue=False,
|
|
246
|
+
show_progress='full').then(
|
|
247
|
+
# update the plot when the router data changes
|
|
248
|
+
fn=update_router_plot,
|
|
249
|
+
inputs=[__fields[__k] for __k in ['position_block', 'router_state']],
|
|
250
|
+
outputs=__fields['plot_block'],
|
|
251
|
+
queue=False,
|
|
252
|
+
show_progress='full')
|
|
253
|
+
# update the range of the position slider when the settings change
|
|
120
254
|
__fields['tokens_block'].change(
|
|
121
255
|
fn=update_position_range,
|
|
122
|
-
inputs=[__fields[__k] for __k in ['position_block', 'tokens_block']],
|
|
256
|
+
inputs=[__fields[__k] for __k in ['position_block', 'tokens_block', 'output_state']],
|
|
123
257
|
outputs=__fields['position_block'],
|
|
124
258
|
queue=False,
|
|
125
259
|
show_progress='hidden')
|
|
260
|
+
# update the plot when the focus changes
|
|
261
|
+
__fields['position_block'].change(
|
|
262
|
+
fn=update_router_plot,
|
|
263
|
+
inputs=[__fields[__k] for __k in ['position_block', 'router_state']],
|
|
264
|
+
outputs=__fields['plot_block'],
|
|
265
|
+
queue=False,
|
|
266
|
+
show_progress='full')
|
|
267
|
+
# update the token highlight when the token focus changes
|
|
268
|
+
__fields['position_block'].change(
|
|
269
|
+
fn=__highlight,
|
|
270
|
+
inputs=[__fields[__k] for __k in ['position_block', 'output_state']],
|
|
271
|
+
outputs=__fields['output_block'],
|
|
272
|
+
queue=False,
|
|
273
|
+
show_progress='hidden')
|
|
126
274
|
# gradio application
|
|
127
275
|
return __app
|
|
128
276
|
|
psaiops/score/router/lib.py
CHANGED
|
@@ -1,66 +1,59 @@
|
|
|
1
1
|
import functools
|
|
2
2
|
|
|
3
3
|
import torch
|
|
4
|
-
import transformers
|
|
5
4
|
|
|
6
|
-
|
|
5
|
+
# COMPUTE ########################################################################
|
|
7
6
|
|
|
8
|
-
|
|
9
|
-
|
|
10
|
-
@functools.lru_cache(maxsize=4)
|
|
11
|
-
def get_tokenizer(name: str, device: str='cpu'):
|
|
12
|
-
return transformers.AutoTokenizer.from_pretrained(
|
|
13
|
-
name,
|
|
14
|
-
use_fast=True,
|
|
15
|
-
dtype='auto',
|
|
16
|
-
device_map=device)
|
|
17
|
-
|
|
18
|
-
@functools.lru_cache(maxsize=2)
|
|
19
|
-
def get_model(name: str, device: str='cpu'):
|
|
20
|
-
__model = deformers.models.openai.gptoss.GptOssForCausalInference.from_pretrained(
|
|
21
|
-
name,
|
|
22
|
-
dtype='auto',
|
|
23
|
-
device_map=device)
|
|
24
|
-
# toggle the inference mode (not training)
|
|
25
|
-
__model.eval()
|
|
26
|
-
# transformers model
|
|
27
|
-
return __model
|
|
28
|
-
|
|
29
|
-
# PREPROCESS #####################################################################
|
|
30
|
-
|
|
31
|
-
@functools.lru_cache(maxsize=4)
|
|
32
|
-
def preprocess_token_ids(
|
|
33
|
-
tokenizer_obj: object,
|
|
34
|
-
prompt_str: str,
|
|
35
|
-
device_str: str='cpu'
|
|
36
|
-
) -> dict:
|
|
37
|
-
# tokenize
|
|
38
|
-
__inputs = tokenizer_obj(prompt_str, return_tensors='pt')
|
|
39
|
-
# move to the main device
|
|
40
|
-
return {__k: __v.to(device_str) for __k, __v in __inputs.items()}
|
|
41
|
-
|
|
42
|
-
# GENERATE #######################################################################
|
|
43
|
-
|
|
44
|
-
def generate_token_ids(
|
|
7
|
+
def compute_router_weights(
|
|
45
8
|
model_obj: object,
|
|
46
|
-
|
|
47
|
-
token_num: int,
|
|
48
|
-
topk_num: int = 4,
|
|
49
|
-
topp_num: float = 0.9,
|
|
9
|
+
token_data: torch.Tensor,
|
|
50
10
|
) -> torch.Tensor:
|
|
51
|
-
#
|
|
11
|
+
# process the full sequence
|
|
52
12
|
with torch.no_grad():
|
|
53
|
-
__outputs = model_obj
|
|
54
|
-
|
|
55
|
-
max_new_tokens=token_num,
|
|
56
|
-
do_sample=(0.0 < topp_num < 1.0) or (topk_num > 0),
|
|
57
|
-
top_k=topk_num if (topk_num > 0) else None,
|
|
58
|
-
top_p=topp_num if (0.0 < topp_num < 1.0) else None,
|
|
59
|
-
return_dict_in_generate=True,
|
|
60
|
-
output_hidden_states=False,
|
|
13
|
+
__outputs = model_obj(
|
|
14
|
+
input_ids=token_data,
|
|
61
15
|
output_attentions=False,
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
16
|
+
output_router_logits=True,
|
|
17
|
+
return_dict=True)
|
|
18
|
+
# stack all the layer outputs L * (T, E) => (L, T, E)
|
|
19
|
+
__logits = torch.stack(__outputs.router_logits, dim=0)
|
|
20
|
+
# turn the logits into expert probabilities
|
|
21
|
+
return torch.softmax(__logits, dim=-1)
|
|
22
|
+
|
|
23
|
+
# REDUCE #######################################################################
|
|
24
|
+
|
|
25
|
+
def reduce_router_weights(
|
|
26
|
+
router_data: torch.Tensor,
|
|
27
|
+
token_idx: int, # -1 => avg over all tokens
|
|
28
|
+
) -> torch.Tensor:
|
|
29
|
+
# parse
|
|
30
|
+
__layer_dim, __token_dim, __expert_dim = tuple(router_data.shape) # L, T, E
|
|
31
|
+
__token_idx = min(token_idx, __token_dim - 1)
|
|
32
|
+
# select the relevant data along each axis
|
|
33
|
+
__token_slice = slice(0, __token_dim) if (__token_idx < 0) else slice(__token_idx, __token_idx + 1)
|
|
34
|
+
# filter the data
|
|
35
|
+
__data = router_data[slice(None), __token_slice, slice(None)]
|
|
36
|
+
# reduce all the axes but the last
|
|
37
|
+
return __data.mean(dim=1, keepdim=False)
|
|
38
|
+
|
|
39
|
+
# FORMAT #########################################################################
|
|
40
|
+
|
|
41
|
+
def postprocess_router_weights(
|
|
42
|
+
router_data: torch.Tensor, # (L, E)
|
|
43
|
+
) -> list:
|
|
44
|
+
# the averaging over tokens may have broken the scaling
|
|
45
|
+
__probs = torch.softmax(router_data, dim=-1)
|
|
46
|
+
# enforce the output range [0; 1] with 1 included
|
|
47
|
+
return __probs / __probs.amax(dim=-1, keepdim=True)
|
|
48
|
+
|
|
49
|
+
# POSTPROCESS ####################################################################
|
|
50
|
+
|
|
51
|
+
def postprocess_token_cls(
|
|
52
|
+
token_idx: int,
|
|
53
|
+
token_dim: int,
|
|
54
|
+
) -> list:
|
|
55
|
+
__token_idx = max(-1, min(token_dim, token_idx))
|
|
56
|
+
# class 1 for the focused token(s) 0 for the rest
|
|
57
|
+
__token_cls = [str(int(__i == token_idx)) for __i in range(token_dim)]
|
|
58
|
+
# average on all the tokens when the idx is negative
|
|
59
|
+
return token_dim * ['1'] if (token_idx < 0) else __token_cls
|
psaiops/score/shapley/app.py
CHANGED
|
@@ -4,12 +4,13 @@ import gradio
|
|
|
4
4
|
import torch
|
|
5
5
|
import torch.cuda
|
|
6
6
|
|
|
7
|
-
import psaiops.
|
|
7
|
+
import psaiops.common.model
|
|
8
|
+
import psaiops.common.tokenizer
|
|
8
9
|
|
|
9
10
|
# META #########################################################################
|
|
10
11
|
|
|
11
12
|
STYLE = '''.white-text span { color: white; }'''
|
|
12
|
-
TITLE = '''
|
|
13
|
+
TITLE = '''Shapley Scoring'''
|
|
13
14
|
INTRO = '''Score each token according to their [Shapley value](https://en.wikipedia.org/wiki/Shapley_value).\nUnder construction, only "openai/gpt-oss-20b" is available for now.'''
|
|
14
15
|
|
|
15
16
|
MODEL = 'openai/gpt-oss-20b'
|
|
@@ -128,8 +129,8 @@ def create_app(title: str=TITLE, intro: str=INTRO, style: str=STYLE, model: str=
|
|
|
128
129
|
with gradio.Blocks(theme=gradio.themes.Soft(), title=title, css=style) as __app:
|
|
129
130
|
# load the model
|
|
130
131
|
__device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
131
|
-
# __model = psaiops.
|
|
132
|
-
__tokenizer = psaiops.
|
|
132
|
+
# __model = psaiops.common.model.get_model(name=model, device=__device)
|
|
133
|
+
__tokenizer = psaiops.common.tokenizer.get_tokenizer(name=model, device=__device)
|
|
133
134
|
# create the UI
|
|
134
135
|
__fields.update(create_layout(intro=intro))
|
|
135
136
|
# init the state
|
psaiops/score/shapley/lib.py
CHANGED
|
@@ -1,66 +1 @@
|
|
|
1
|
-
import functools
|
|
2
|
-
|
|
3
1
|
import torch
|
|
4
|
-
import transformers
|
|
5
|
-
|
|
6
|
-
import deformers.models.openai.gptoss
|
|
7
|
-
|
|
8
|
-
# LOAD #########################################################################
|
|
9
|
-
|
|
10
|
-
@functools.lru_cache(maxsize=4)
|
|
11
|
-
def get_tokenizer(name: str, device: str='cpu'):
|
|
12
|
-
return transformers.AutoTokenizer.from_pretrained(
|
|
13
|
-
name,
|
|
14
|
-
use_fast=True,
|
|
15
|
-
dtype='auto',
|
|
16
|
-
device_map=device)
|
|
17
|
-
|
|
18
|
-
@functools.lru_cache(maxsize=2)
|
|
19
|
-
def get_model(name: str, device: str='cpu'):
|
|
20
|
-
__model = deformers.models.openai.gptoss.GptOssForCausalInference.from_pretrained(
|
|
21
|
-
name,
|
|
22
|
-
dtype='auto',
|
|
23
|
-
device_map=device)
|
|
24
|
-
# toggle the inference mode (not training)
|
|
25
|
-
__model.eval()
|
|
26
|
-
# transformers model
|
|
27
|
-
return __model
|
|
28
|
-
|
|
29
|
-
# PREPROCESS #####################################################################
|
|
30
|
-
|
|
31
|
-
@functools.lru_cache(maxsize=4)
|
|
32
|
-
def preprocess_token_ids(
|
|
33
|
-
tokenizer_obj: object,
|
|
34
|
-
prompt_str: str,
|
|
35
|
-
device_str: str='cpu'
|
|
36
|
-
) -> dict:
|
|
37
|
-
# tokenize
|
|
38
|
-
__inputs = tokenizer_obj(prompt_str, return_tensors='pt')
|
|
39
|
-
# move to the main device
|
|
40
|
-
return {__k: __v.to(device_str) for __k, __v in __inputs.items()}
|
|
41
|
-
|
|
42
|
-
# GENERATE #######################################################################
|
|
43
|
-
|
|
44
|
-
def generate_token_ids(
|
|
45
|
-
model_obj: object,
|
|
46
|
-
input_args: dict,
|
|
47
|
-
token_num: int,
|
|
48
|
-
topk_num: int = 4,
|
|
49
|
-
topp_num: float = 0.9,
|
|
50
|
-
) -> torch.Tensor:
|
|
51
|
-
# generate completion
|
|
52
|
-
with torch.no_grad():
|
|
53
|
-
__outputs = model_obj.generate(
|
|
54
|
-
**input_args,
|
|
55
|
-
max_new_tokens=token_num,
|
|
56
|
-
do_sample=(0.0 < topp_num < 1.0) or (topk_num > 0),
|
|
57
|
-
top_k=topk_num if (topk_num > 0) else None,
|
|
58
|
-
top_p=topp_num if (0.0 < topp_num < 1.0) else None,
|
|
59
|
-
return_dict_in_generate=True,
|
|
60
|
-
output_hidden_states=False,
|
|
61
|
-
output_attentions=False,
|
|
62
|
-
output_scores=False,
|
|
63
|
-
# early_stopping=True,
|
|
64
|
-
use_cache=True)
|
|
65
|
-
# full sequence
|
|
66
|
-
return __outputs.sequences # (1, T)
|
|
File without changes
|
|
@@ -0,0 +1,152 @@
|
|
|
1
|
+
import functools
|
|
2
|
+
|
|
3
|
+
import gradio
|
|
4
|
+
import torch
|
|
5
|
+
import torch.cuda
|
|
6
|
+
|
|
7
|
+
import psaiops.common.model
|
|
8
|
+
import psaiops.common.tokenizer
|
|
9
|
+
|
|
10
|
+
# META #########################################################################
|
|
11
|
+
|
|
12
|
+
STYLE = '''.white-text span { color: white; }'''
|
|
13
|
+
TITLE = '''Similarity Scoring'''
|
|
14
|
+
INTRO = '''Score each token according to its similarity with a reference prompt.\nUnder construction, only "openai/gpt-oss-20b" is available for now.'''
|
|
15
|
+
|
|
16
|
+
MODEL = 'openai/gpt-oss-20b'
|
|
17
|
+
|
|
18
|
+
# COLORS #######################################################################
|
|
19
|
+
|
|
20
|
+
def create_color_map() -> dict:
|
|
21
|
+
return {
|
|
22
|
+
'-1': '#004444',
|
|
23
|
+
**{str(__i): '#{:02x}0000'.format(int(2.55 * __i)) for __i in range(101)}}
|
|
24
|
+
|
|
25
|
+
# INTRO ########################################################################
|
|
26
|
+
|
|
27
|
+
def create_intro_block(intro: str) -> dict:
|
|
28
|
+
__intro = gradio.Markdown(intro, line_breaks=True)
|
|
29
|
+
return {'intro_block': __intro}
|
|
30
|
+
|
|
31
|
+
# MODEL ########################################################################
|
|
32
|
+
|
|
33
|
+
def create_model_block() -> dict:
|
|
34
|
+
__model = gradio.Dropdown(label='Model', value='openai/gpt-oss-20b', choices=['openai/gpt-oss-20b'], scale=1, allow_custom_value=False, multiselect=False, interactive=True) # 'openai/gpt-oss-120b'
|
|
35
|
+
return {'model_block': __model,}
|
|
36
|
+
|
|
37
|
+
# SAMPLING #####################################################################
|
|
38
|
+
|
|
39
|
+
def create_sampling_block() -> dict:
|
|
40
|
+
__tokens = gradio.Slider(label='Tokens', value=16, minimum=1, maximum=128, step=1, scale=1, interactive=True)
|
|
41
|
+
__topk = gradio.Slider(label='Top K', value=4, minimum=1, maximum=8, step=1, scale=1, interactive=True)
|
|
42
|
+
__topp = gradio.Slider(label='Top P', value=0.9, minimum=0.0, maximum=1.0, step=0.1, scale=1, interactive=True)
|
|
43
|
+
return {
|
|
44
|
+
'tokens_block': __tokens,
|
|
45
|
+
'topk_block': __topk,
|
|
46
|
+
'topp_block': __topp,}
|
|
47
|
+
|
|
48
|
+
# INPUTS #######################################################################
|
|
49
|
+
|
|
50
|
+
def create_inputs_block() -> dict:
|
|
51
|
+
__reference = gradio.Textbox(label='Reference', value='', placeholder='The text used as reference for the comparison.', lines=4, scale=1, show_copy_button=True, interactive=True)
|
|
52
|
+
__prompt = gradio.Textbox(label='Prompt', value='', placeholder='A string of tokens to compare to the reference.', lines=4, scale=1, show_copy_button=True, interactive=True)
|
|
53
|
+
return {
|
|
54
|
+
'reference_input_block': __reference,
|
|
55
|
+
'prompt_input_block': __prompt,}
|
|
56
|
+
|
|
57
|
+
# OUTPUTS ######################################################################
|
|
58
|
+
|
|
59
|
+
def create_outputs_block() -> dict:
|
|
60
|
+
__reference = gradio.HighlightedText(label='', value='', scale=1, interactive=False, show_label=False, show_legend=False, show_inline_category=False, combine_adjacent=False, color_map=create_color_map(), elem_classes='white-text')
|
|
61
|
+
__prompt = gradio.HighlightedText(label='', value='', scale=1, interactive=False, show_label=False, show_legend=False, show_inline_category=False, combine_adjacent=False, color_map=create_color_map(), elem_classes='white-text')
|
|
62
|
+
return {
|
|
63
|
+
'reference_output_block': __reference,
|
|
64
|
+
'prompt_output_block': __prompt,}
|
|
65
|
+
|
|
66
|
+
# SELECT #######################################################################
|
|
67
|
+
|
|
68
|
+
def create_selection_block() -> dict:
|
|
69
|
+
__position = gradio.Slider(label='Token Position', value=-1, minimum=-1, maximum=15, step=1, scale=1, interactive=True) # info='-1 to average on all tokens'
|
|
70
|
+
__layer = gradio.Slider(label='Layer Depth', value=12, minimum=-1, maximum=23, step=1, scale=1, interactive=True) # info='-1 to average on all layers'
|
|
71
|
+
return {
|
|
72
|
+
'position_block': __position,
|
|
73
|
+
'layer_block': __layer,}
|
|
74
|
+
|
|
75
|
+
# ACTIONS ######################################################################
|
|
76
|
+
|
|
77
|
+
def create_actions_block() -> dict:
|
|
78
|
+
__process = gradio.Button('Process', variant='primary', size='lg', scale=1, interactive=True)
|
|
79
|
+
return {'process_block': __process,}
|
|
80
|
+
|
|
81
|
+
# STATE ########################################################################
|
|
82
|
+
|
|
83
|
+
def create_state() -> dict:
|
|
84
|
+
return {}
|
|
85
|
+
|
|
86
|
+
# LAYOUT #######################################################################
|
|
87
|
+
|
|
88
|
+
def create_layout(intro: str=INTRO) -> dict:
|
|
89
|
+
__fields = {}
|
|
90
|
+
__fields.update(create_intro_block(intro=intro))
|
|
91
|
+
with gradio.Tabs():
|
|
92
|
+
with gradio.Tab('Score Tokens') as __main_tab:
|
|
93
|
+
__fields.update({'main_tab': __main_tab})
|
|
94
|
+
with gradio.Row(equal_height=True):
|
|
95
|
+
__fields.update(create_inputs_block())
|
|
96
|
+
with gradio.Row(equal_height=True):
|
|
97
|
+
__fields.update(create_outputs_block())
|
|
98
|
+
with gradio.Row(equal_height=True):
|
|
99
|
+
__fields.update(create_selection_block())
|
|
100
|
+
with gradio.Row(equal_height=True):
|
|
101
|
+
__fields.update(create_actions_block())
|
|
102
|
+
with gradio.Tab('Settings') as __settings_tab:
|
|
103
|
+
__fields.update({'settings_tab': __settings_tab})
|
|
104
|
+
with gradio.Column(scale=1):
|
|
105
|
+
with gradio.Row(equal_height=True):
|
|
106
|
+
__fields.update(create_model_block())
|
|
107
|
+
with gradio.Row(equal_height=True):
|
|
108
|
+
__fields.update(create_sampling_block())
|
|
109
|
+
return __fields
|
|
110
|
+
|
|
111
|
+
# EVENTS #######################################################################
|
|
112
|
+
|
|
113
|
+
def update_layer_range(value: float, model: str) -> dict:
|
|
114
|
+
return gradio.update(maximum=35, value=min(35, int(value))) if '120b' in model else gradio.update(maximum=23, value=min(23, int(value)))
|
|
115
|
+
|
|
116
|
+
def update_position_range(value: float, tokens: float) -> dict:
|
|
117
|
+
return gradio.update(maximum=int(tokens) - 1, value=min(int(tokens) - 1, int(value)))
|
|
118
|
+
|
|
119
|
+
# APP ##########################################################################
|
|
120
|
+
|
|
121
|
+
def create_app(title: str=TITLE, intro: str=INTRO, style: str=STYLE, model: str=MODEL) -> gradio.Blocks:
|
|
122
|
+
__fields = {}
|
|
123
|
+
with gradio.Blocks(theme=gradio.themes.Soft(), title=title, css=style) as __app:
|
|
124
|
+
# load the model
|
|
125
|
+
__device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
|
126
|
+
# __model = psaiops.common.model.get_model(name=model, device=__device)
|
|
127
|
+
__tokenizer = psaiops.common.tokenizer.get_tokenizer(name=model, device=__device)
|
|
128
|
+
# create the UI
|
|
129
|
+
__fields.update(create_layout(intro=intro))
|
|
130
|
+
# init the state
|
|
131
|
+
__fields.update(create_state())
|
|
132
|
+
# wire the input fields
|
|
133
|
+
__fields['tokens_block'].change(
|
|
134
|
+
fn=update_position_range,
|
|
135
|
+
inputs=[__fields[__k] for __k in ['position_block', 'tokens_block']],
|
|
136
|
+
outputs=__fields['position_block'],
|
|
137
|
+
queue=False,
|
|
138
|
+
show_progress='hidden')
|
|
139
|
+
__fields['model_block'].change(
|
|
140
|
+
fn=update_layer_range,
|
|
141
|
+
inputs=[__fields[__k] for __k in ['layer_block', 'model_block']],
|
|
142
|
+
outputs=__fields['layer_block'],
|
|
143
|
+
queue=False,
|
|
144
|
+
show_progress='hidden')
|
|
145
|
+
# gradio application
|
|
146
|
+
return __app
|
|
147
|
+
|
|
148
|
+
# MAIN #########################################################################
|
|
149
|
+
|
|
150
|
+
if __name__ == '__main__':
|
|
151
|
+
__app = create_app()
|
|
152
|
+
__app.launch(share=True, debug=True)
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
import torch
|
|
@@ -1,23 +1,19 @@
|
|
|
1
|
-
Metadata-Version: 2.
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
2
|
Name: psaiops
|
|
3
|
-
Version: 0.
|
|
3
|
+
Version: 0.4.0
|
|
4
4
|
Summary: Web apps to inspect & engineer NN activations.
|
|
5
|
-
|
|
6
|
-
|
|
7
|
-
|
|
8
|
-
Requires-Python:
|
|
9
|
-
|
|
10
|
-
|
|
11
|
-
|
|
12
|
-
|
|
13
|
-
|
|
14
|
-
|
|
15
|
-
Requires-Dist:
|
|
16
|
-
Requires-Dist:
|
|
17
|
-
Requires-Dist: gradio (>=5.0)
|
|
18
|
-
Requires-Dist: kernels (==0.10)
|
|
19
|
-
Requires-Dist: requests (>=2.0)
|
|
20
|
-
Requires-Dist: triton (==3.4)
|
|
5
|
+
Author-email: apehex <17317183+apehex@users.noreply.github.com>
|
|
6
|
+
License-Expression: AGPL-3.0
|
|
7
|
+
License-File: .github/LICENSE.md
|
|
8
|
+
Requires-Python: <3.14,>=3.10
|
|
9
|
+
Requires-Dist: accelerate>=1.10
|
|
10
|
+
Requires-Dist: deformers>=0.0
|
|
11
|
+
Requires-Dist: gradio>=5.0
|
|
12
|
+
Requires-Dist: kernels==0.10
|
|
13
|
+
Requires-Dist: matplotlib>=3.10
|
|
14
|
+
Requires-Dist: mlable-torch>=0.2
|
|
15
|
+
Requires-Dist: requests>=2.0
|
|
16
|
+
Requires-Dist: triton==3.4
|
|
21
17
|
Description-Content-Type: text/markdown
|
|
22
18
|
|
|
23
19
|
# psAI ops <img src="images/logo.png" alt="apehex logo" width="32" height="32">
|
|
@@ -36,4 +32,3 @@ Licensed under the [aGPLv3][github-license].
|
|
|
36
32
|
|
|
37
33
|
[shield-license]: https://img.shields.io/badge/license-aGPLv3-green?style=flat-square
|
|
38
34
|
[shield-release]: https://img.shields.io/github/release/apehex/psai-ops.svg?style=flat-square
|
|
39
|
-
|