psaiops 0.0.13__py3-none-any.whl → 0.4.0__py3-none-any.whl

psaiops/score/residual/app.py

@@ -0,0 +1,290 @@
+import functools
+
+import gradio
+import torch
+import torch.cuda
+import matplotlib.pyplot
+
+import psaiops.common.model
+import psaiops.common.tokenizer
+import psaiops.score.residual.lib
+
+# META #########################################################################
+
+STYLE = '''.white-text span { color: white; }'''
+TITLE = '''Router Scoring'''
+INTRO = '''Plot the logits of the router for a given prompt.\nUnder construction, only "openai/gpt-oss-20b" is available for now.'''
+
+MODEL = 'openai/gpt-oss-20b'
+
+# COLORS #######################################################################
+
+def create_color_map() -> dict:
+    return {
+        '0': '#000000',
+        '1': '#004444',}
+
+# INTRO ########################################################################
+
+def create_intro_block(intro: str) -> dict:
+    __intro = gradio.Markdown(intro, line_breaks=True)
+    return {'intro_block': __intro}
+
+# MODEL ########################################################################
+
+def create_model_block() -> dict:
+    __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'
+    return {'model_block': __model,}
+
+# SAMPLING #####################################################################
+
+def create_sampling_block() -> dict:
+    __tokens = gradio.Slider(label='Tokens', value=16, minimum=1, maximum=128, step=1, scale=1, interactive=True)
+    __topk = gradio.Slider(label='Top K', value=4, minimum=1, maximum=8, step=1, scale=1, interactive=True)
+    __topp = gradio.Slider(label='Top P', value=0.9, minimum=0.0, maximum=1.0, step=0.1, scale=1, interactive=True)
+    return {
+        'tokens_block': __tokens,
+        'topk_block': __topk,
+        'topp_block': __topp,}
+
+# INPUTS #######################################################################
+
+def create_inputs_block() -> dict:
+    __input = gradio.Textbox(label='Prompt', value='', placeholder='A string of tokens to score.', lines=4, scale=1, show_copy_button=True, interactive=True)
+    return {'input_block': __input}
+
+# PLOTS ########################################################################
+
+def create_plot_block() -> dict:
+    __plot = gradio.Plot(label='Router', scale=1)
+    return {'plot_block': __plot,}
+
+# OUTPUTS ######################################################################
+
+def create_outputs_block() -> dict:
+    __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')
+    return {'output_block': __output}
+
+# SELECT #######################################################################
+
+def create_selection_block() -> dict:
+    # __play = gradio.Button('>', variant='primary', size='lg', scale=1, interactive=True)
+    __position = gradio.Slider(label='Token', value=-1, minimum=-1, maximum=15, step=1, scale=1, interactive=True) # info='-1 to average on all tokens'
+    return {'position_block': __position,}
+
+# ACTIONS ######################################################################
+
+def create_actions_block() -> dict:
+    __process = gradio.Button('Process', variant='primary', size='lg', scale=1, interactive=True)
+    return {'process_block': __process,}
+
+# STATE ########################################################################
+
+def create_state() -> dict:
+    return {
+        'output_state': gradio.State(None),
+        'hidden_state': gradio.State(None),}
+
+# LAYOUT #######################################################################
+
+def create_layout(intro: str=INTRO) -> dict:
+    __fields = {}
+    __fields.update(create_intro_block(intro=intro))
+    with gradio.Tabs():
+        with gradio.Tab('Score Tokens') as __main_tab:
+            __fields.update({'main_tab': __main_tab})
+            with gradio.Row(equal_height=True):
+                __fields.update(create_inputs_block())
+            with gradio.Row(equal_height=True):
+                __fields.update(create_plot_block())
+            with gradio.Row(equal_height=True):
+                __fields.update(create_outputs_block())
+            with gradio.Row(equal_height=True):
+                __fields.update(create_selection_block())
+            with gradio.Row(equal_height=True):
+                __fields.update(create_actions_block())
+        with gradio.Tab('Settings') as __settings_tab:
+            __fields.update({'settings_tab': __settings_tab})
+            with gradio.Column(scale=1):
+                with gradio.Row(equal_height=True):
+                    __fields.update(create_model_block())
+                with gradio.Row(equal_height=True):
+                    __fields.update(create_sampling_block())
+    return __fields
+
+# EVENTS #######################################################################
+
+def update_position_range(
+    current_val: float,
+    token_num: float,
+    output_data: torch.Tensor,
+) -> dict:
+    # take the generated tokens into account
+    __max = int(token_num) - 1 if (output_data is None) else int(output_data.shape[-1])
+    # keep the previous value if possible
+    __val = min(int(current_val), __max)
+    # return a gradio update dictionary
+    return gradio.update(maximum=__max, value=__val)
+
+def update_computation_state(
+    token_num: float,
+    topk_num: float,
+    topp_num: float,
+    token_idx: float,
+    prompt_str: str,
+    device_str: str,
+    model_obj: object,
+    tokenizer_obj: object,
+) -> tuple:
+    # sanitize the inputs
+    __token_num = max(1, min(128, int(token_num)))
+    __topk_num = max(1, min(8, int(topk_num)))
+    __topp_num = max(0.0, min(1.0, float(topp_num)))
+    __token_idx = max(-1, min(__token_num, int(token_idx)))
+    __prompt_str = prompt_str.strip()
+    __device_str = device_str if (device_str in ['cpu', 'cuda']) else 'cpu'
+    # exit if some values are missing
+    if (not __prompt_str) or (model_obj is None) or (tokenizer_obj is None):
+        return (torch.empty(0), torch.empty(0))
+    # dictionary {'input_ids': _, 'attention_mask': _}
+    __input_data = psaiops.common.tokenizer.preprocess_token_ids(
+        tokenizer_obj=tokenizer_obj,
+        prompt_str=__prompt_str,
+        device_str=__device_str)
+    # tensor (1, T) and O * L * (1, I, H)
+    __output_data, __hidden_data = psaiops.score.residual.lib.generate_token_ids(
+        model_obj=model_obj,
+        input_args=__input_data,
+        token_num=__token_num,
+        topk_num=__topk_num,
+        topp_num=__topp_num)
+    # tensor (1, L, I + O, H)
+    __hidden_data = psaiops.score.residual.lib.merge_hidden_states(
+        hidden_data=__hidden_data)
+    # update each component => (highlight, plot) states
+    return (
+        __output_data.cpu().float(),
+        __hidden_data.cpu().float(),)
+
+def update_hidden_plot(
+    token_idx: float,
+    hidden_data: torch.Tensor,
+) -> tuple:
+    # exit if some values are missing
+    if (hidden_data is None) or (len(hidden_data) == 0):
+        return None
+    # reduce the token axis (B, L, T, E) => (B, L, E)
+    __plot_data = psaiops.score.residual.lib.reduce_hidden_states(
+        hidden_data=hidden_data,
+        token_idx=int(token_idx),)
+    # rescale the data to [-1; 1] (B, L, E)
+    __plot_data = psaiops.score.residual.lib.rescale_hidden_states(
+        hidden_data=__plot_data)
+    # reshape into a 3D tensor by folding E (B, L, E) => (B, W, H, L)
+    __plot_data = psaiops.score.residual.lib.reshape_hidden_states(
+        hidden_data=__plot_data)
+    # map the [-1; 1] activations to RGBA colors
+    __plot_data = psaiops.score.residual.lib.color_hidden_states(
+        hidden_data=__plot_data.numpy())
+    # mask the small activations to improve the plot readability
+    __mask_data = psaiops.score.residual.lib.mask_hidden_states(
+        hidden_data=__plot_data,
+        topk_num=128).numpy()
+    # plot the first sample
+    __figure, __axes = matplotlib.pyplot.subplots(111, projection='3d')
+    __axes.voxels(filled=__mask_data[0].numpy(), facecolors=__plot_data[0], edgecolor=None)
+    __figure.tight_layout()
+    # remove the figure for the pyplot register for garbage collection
+    matplotlib.pyplot.close(__figure)
+    # update each component => (highlight, plot) states
+    return __figure
+
+def update_text_highlight(
+    token_idx: float,
+    output_data: torch.Tensor,
+    tokenizer_obj: object,
+) -> list:
+    # exit if some values are missing
+    if (output_data is None) or (len(output_data) == 0):
+        return None
+    # detokenize the IDs
+    __token_str = psaiops.common.tokenizer.postprocess_token_ids(
+        tokenizer_obj=tokenizer_obj,
+        token_data=output_data)
+    # list of string classes
+    __token_cls = psaiops.score.residual.lib.postprocess_token_cls(
+        token_idx=int(token_idx),
+        token_dim=len(__token_str))
+    # pairs of token and class
+    return list(zip(__token_str, __token_cls))
+
+# APP ##########################################################################
+
+def create_app(title: str=TITLE, intro: str=INTRO, style: str=STYLE, model: str=MODEL) -> gradio.Blocks:
+    __fields = {}
+    with gradio.Blocks(theme=gradio.themes.Soft(), title=title, css=style) as __app:
+        # load the model
+        __device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        __model = psaiops.common.model.get_model(name=model, device=__device)
+        __tokenizer = psaiops.common.tokenizer.get_tokenizer(name=model, device=__device)
+        # adapt the event handlers
+        __compute = functools.partial(update_computation_state, model_obj=__model, tokenizer_obj=__tokenizer, device_str=__device)
+        __highlight = functools.partial(update_text_highlight, tokenizer_obj=__tokenizer)
+        # create the UI
+        __fields.update(create_layout(intro=intro))
+        # init the state
+        __fields.update(create_state())
+        # update the data after clicking process
+        __fields['process_block'].click(
+            fn=__compute,
+            inputs=[__fields[__k] for __k in ['tokens_block', 'topk_block', 'topp_block', 'position_block', 'input_block']],
+            outputs=[__fields[__k] for __k in ['output_state', 'hidden_state']],
+            queue=False,
+            show_progress='full').then(
+        # update the range of the position slider when the output changes
+            fn=update_position_range,
+            inputs=[__fields[__k] for __k in ['position_block', 'tokens_block', 'output_state']],
+            outputs=__fields['position_block'],
+            queue=False,
+            show_progress='hidden').then(
+        # update the token highlight when the output data changes
+            fn=__highlight,
+            inputs=[__fields[__k] for __k in ['position_block', 'output_state']],
+            outputs=__fields['output_block'],
+            queue=False,
+            show_progress='full').then(
+        # update the plot when the router data changes
+            fn=update_hidden_plot,
+            inputs=[__fields[__k] for __k in ['position_block', 'hidden_state']],
+            outputs=__fields['plot_block'],
+            queue=False,
+            show_progress='full')
+        # update the range of the position slider when the settings change
+        __fields['tokens_block'].change(
+            fn=update_position_range,
+            inputs=[__fields[__k] for __k in ['position_block', 'tokens_block', 'output_state']],
+            outputs=__fields['position_block'],
+            queue=False,
+            show_progress='hidden')
+        # update the plot when the focus changes
+        __fields['position_block'].change(
+            fn=update_hidden_plot,
+            inputs=[__fields[__k] for __k in ['position_block', 'hidden_state']],
+            outputs=__fields['plot_block'],
+            queue=False,
+            show_progress='full')
+        # update the token highlight when the token focus changes
+        __fields['position_block'].change(
+            fn=__highlight,
+            inputs=[__fields[__k] for __k in ['position_block', 'output_state']],
+            outputs=__fields['output_block'],
+            queue=False,
+            show_progress='hidden')
+        # gradio application
+        return __app
+
+# MAIN #########################################################################
+
+if __name__ == '__main__':
+    __app = create_app()
+    __app.launch(share=True, debug=True)

psaiops/score/residual/lib.py

@@ -0,0 +1,134 @@
+import functools
+import math
+
+import matplotlib
+import numpy
+import torch
+
+# GENERATE #######################################################################
+
+@functools.lru_cache(maxsize=32)
+def generate_token_ids(
+    model_obj: object,
+    input_args: dict,
+    token_num: int,
+    topk_num: int = 4,
+    topp_num: float = 0.9,
+) -> tuple:
+    # generate completion
+    with torch.no_grad():
+        __outputs = model_obj.generate(
+            **input_args,
+            max_new_tokens=token_num,
+            do_sample=(0.0 < topp_num < 1.0) or (topk_num > 0),
+            top_k=topk_num if (topk_num > 0) else None,
+            top_p=topp_num if (0.0 < topp_num < 1.0) else None,
+            return_dict_in_generate=True,
+            output_hidden_states=True,
+            output_attentions=False,
+            output_scores=False,
+            early_stopping=True,
+            use_cache=True)
+    # ((B, T), O * L * (B, I, H))
+    return __outputs.sequences, __outputs.hidden_states
+
+# MERGE ########################################################################
+
+def merge_hidden_states(
+    hidden_data: torch.Tensor,
+) -> torch.Tensor:
+    # parse the inputs
+    __token_dim = len(hidden_data)
+    __layer_dim = len(hidden_data[0])
+    # stack the data for each layer => (B, L, I + O, H)
+    return torch.stack(
+        [
+            # concatenate the data for all the tokens => (B, I + O, H)
+            torch.concatenate([hidden_data[__t][__l] for __t in range(__token_dim)], dim=1)
+            for __l in range(__layer_dim)],
+        dim=1)
+
+# REDUCE #######################################################################
+
+def reduce_hidden_states(
+    hidden_data: torch.Tensor,
+    token_idx: int, # -1 => avg over all tokens
+) -> torch.Tensor:
+    # parse the hidden states (B, L, T, H)
+    __batch_dim, __layer_dim, __token_dim, __hidden_dim = tuple(hidden_data.shape)
+    __token_idx = min(token_idx, __token_dim - 1)
+    # select the relevant data along each axis
+    __token_slice = slice(0, __token_dim) if (__token_idx < 0) else slice(__token_idx, __token_idx + 1)
+    # filter the data
+    __data = hidden_data[slice(None), slice(None), __token_slice, slice(None)]
+    # reduce the token axis => (B, L, H)
+    return __data.mean(dim=2, keepdim=False)
+
+# RESCALE ######################################################################
+
+def rescale_hidden_states(
+    hidden_data: torch.Tensor, # (B, L, H)
+) -> torch.Tensor:
+    # compute the scale of the data, layer by layer
+    __s = torch.quantile(hidden_data.abs(), q=0.9, dim=-1, keepdim=True)
+    # log scaling on large values and linear near 0
+    __a = torch.asinh(hidden_data / (__s + torch.finfo().eps))
+    # clip and map to [-1; 1]
+    return 0.33 * __a.clamp(min=-3, max=3)
+
+# RESHAPE ######################################################################
+
+def reshape_hidden_states(
+    hidden_data: torch.Tensor, # (B, L, H)
+) -> torch.Tensor:
+    # parse the hidden states (B, L, H)
+    __batch_dim, __layer_dim, __hidden_dim = tuple(hidden_data.shape)
+    # factor the hidden dimension
+    __width_dim = math.gcd(__hidden_dim, 2 ** int(math.log2(__hidden_dim))) # greatest power of 2 that divides H
+    __height_dim = __hidden_dim // __width_dim
+    # reshape into (B, W, H, L)
+    return hidden_data.reshape((__batch_dim, __layer_dim, __width_dim, __height_dim)).permute(0, 2, 3, 1)
+
+# MASK #########################################################################
+
+def mask_hidden_states(
+    hidden_data: torch.Tensor, # (B, L, H)
+    topk_num: int=128,
+) -> torch.Tensor:
+    # sanitize
+    __k = min(topk_num, int(hidden_data.shape[-1]))
+    # indices of the topk values
+    __indices = hidden_data.abs().topk(__k, dim=-1, largest=True, sorted=False).indices
+    # initialize the mask with False
+    __mask = torch.zeros_like(hidden_data, dtype=torch.bool)
+    # (B, L, H) mask of the topk values
+    return __mask.scatter_(dim=-1, index=__indices, src=True)
+
+# FORMAT #######################################################################
+
+def color_hidden_states(
+    hidden_data: numpy.array, # (B, H, W, L)
+    gamma_val: float=0.7,
+    alpha_val: float=0.35,
+    color_map: matplotlib.colormaps['coolwarm'],
+) -> list:
+    # [-1; 1] => [0; 1]
+    __data = 0.5 * (hidden_data + 1.0)
+    # (B, W, H, L) => (B, W, H, L, 4)
+    __rgba = color_map[__data]
+    # compute the transparency from the magnitude
+    __rgba[..., 3] = alpha_val * (np.abs(hidden_data) ** gamma_val)
+    # (B, W, H, L, 4) in [0; 1]
+    return __rgba
+
+# POSTPROCESS ##################################################################
+
+def postprocess_token_cls(
+    token_idx: int,
+    token_dim: int,
+) -> list:
+    __token_idx = max(-1, min(token_dim, token_idx))
+    # class 1 for the focused token(s) 0 for the rest
+    __token_cls = [str(int(__i == token_idx)) for __i in range(token_dim)]
+    # average on all the tokens when the idx is negative
+    return token_dim * ['1'] if (token_idx < 0) else __token_cls