psaiops 0.3.2__py3-none-any.whl → 0.4.2__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.
@@ -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 = '''Attention Scoring'''
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
- '-1': '#004444',
22
- **{str(__i): '#{:02x}0000'.format(int(2.55 * __i)) for __i in range(101)}}
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.Plot(label='Logits', scale=1)
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
- __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'
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(value: float, tokens: float) -> dict:
104
- return gradio.update(maximum=int(tokens) - 1, value=min(int(tokens) - 1, int(value)))
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_ids=__input_data['input_ids'],
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
- # __model = psaiops.score.router.lib.get_model(name=model, device=__device)
114
- __tokenizer = psaiops.score.router.lib.get_tokenizer(name=model, device=__device)
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
- # wire the input fields
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
 
@@ -1,66 +1,59 @@
1
1
  import functools
2
2
 
3
3
  import torch
4
- import transformers
5
4
 
6
- import deformers.models.openai.gptoss
5
+ # COMPUTE ########################################################################
7
6
 
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(
7
+ def compute_router_weights(
45
8
  model_obj: object,
46
- input_args: dict,
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
- # generate completion
11
+ # process the full sequence
52
12
  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,
13
+ __outputs = model_obj(
14
+ input_ids=token_data,
61
15
  output_attentions=False,
62
- output_scores=False,
63
- # early_stopping=True,
64
- use_cache=True)
65
- # full sequence
66
- return __outputs.sequences # (1, T)
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
@@ -4,12 +4,13 @@ import gradio
4
4
  import torch
5
5
  import torch.cuda
6
6
 
7
- import psaiops.score.shapley.lib
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 = '''Attention Scoring'''
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.score.shapley.lib.get_model(name=model, device=__device)
132
- __tokenizer = psaiops.score.shapley.lib.get_tokenizer(name=model, device=__device)
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
@@ -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.3
1
+ Metadata-Version: 2.4
2
2
  Name: psaiops
3
- Version: 0.3.2
3
+ Version: 0.4.2
4
4
  Summary: Web apps to inspect & engineer NN activations.
5
- License: .github/LICENSE.md
6
- Author: apehex
7
- Author-email: apehex@protonmail.com
8
- Requires-Python: >=3.10,<3.14
9
- Classifier: License :: Other/Proprietary License
10
- Classifier: Programming Language :: Python :: 3
11
- Classifier: Programming Language :: Python :: 3.10
12
- Classifier: Programming Language :: Python :: 3.11
13
- Classifier: Programming Language :: Python :: 3.12
14
- Classifier: Programming Language :: Python :: 3.13
15
- Requires-Dist: accelerate (>=1.10)
16
- Requires-Dist: deformers (>=0.0)
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
-