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

psaiops/score/attention/app.py

@@ -4,13 +4,15 @@ import gradio
 import torch
 import torch.cuda
 
+import psaiops.common.model
+import psaiops.common.tokenizer
 import psaiops.score.attention.lib
 
 # META #########################################################################
 
-TITLE = '''Attention Scoring'''
-INTRO = '''Score each token according to the weights of the attention layers.\nThe model is fixed to "openai/gpt-oss-20b" for now.'''
 STYLE = '''.white-text span { color: white; }'''
+TITLE = '''Attention Scoring'''
+INTRO = '''Score each token according to the weights of the attention layers.\nUnder construction, only "openai/gpt-oss-20b" is available for now.'''
 
 MODEL = 'openai/gpt-oss-20b'
 
@@ -18,26 +20,20 @@ MODEL = 'openai/gpt-oss-20b'
 
 def create_color_map() -> dict:
     return {
-        '-1': '#00ff00',
+        '-1': '#004444',
         **{str(__i): '#{:02x}0000'.format(int(2.55 * __i)) for __i in range(101)}}
 
 # INTRO ########################################################################
 
 def create_intro_block(intro: str) -> dict:
-    __intro = gradio.Markdown(intro)
+    __intro = gradio.Markdown(intro, line_breaks=True)
     return {'intro_block': __intro}
 
 # MODEL ########################################################################
 
 def create_model_block() -> dict:
-    __model_dd = 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'
-    __layer_sl = gradio.Slider(label='Layer Depth', value=12, minimum=-1, maximum=23, step=1, scale=1, interactive=True) # info='-1 to average on all layers'
-    __head_sl = gradio.Slider(label='Attention Head', value=-1, minimum=-1, maximum=63, step=1, scale=1, interactive=True) # info='-1 to average on all heads'
-    __model_dd.change(fn=update_layer_range, inputs=[__layer_sl, __model_dd], outputs=__layer_sl, queue=False, show_progress='hidden')
-    return {
-        'model_block': __model_dd,
-        'layer_block': __layer_sl,
-        'head_block': __head_sl}
+    __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 #####################################################################
 
@@ -53,14 +49,14 @@ def create_sampling_block() -> dict:
 # TARGET #######################################################################
 
 def create_target_block() -> dict:
-    __target = gradio.Radio(label='Score', value='Inputs', choices=['Inputs', 'Outputs', 'Both'], scale=1, interactive=True)
+    __target = gradio.Radio(label='Score', value='Inputs', choices=['Inputs', 'Everything'], scale=1, interactive=True)
     return {'target_block': __target}
 
 # DISPLAY ######################################################################
 
-def create_display_block() -> dict:
-    __display = gradio.Radio(label='Display', value='Tokens', choices=['Tokens', 'Indexes'], scale=1, interactive=True)
-    return {'display_block': __display}
+# def create_display_block() -> dict:
+#     __display = gradio.Radio(label='Display', value='Tokens', choices=['Tokens', 'Indexes'], scale=1, interactive=True)
+#     return {'display_block': __display}
 
 # INPUTS #######################################################################
 
@@ -71,17 +67,25 @@ def create_inputs_block() -> dict:
 # OUTPUTS ######################################################################
 
 def create_outputs_block() -> dict:
-    __output = gradio.HighlightedText(label='Scores', value='', scale=1, interactive=False, show_legend=False, show_inline_category=False, combine_adjacent=True, color_map=create_color_map(), elem_classes='white-text')
+    __output = gradio.HighlightedText(label='Scores', 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:
+    __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'
+    __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'
+    __head = gradio.Slider(label='Attention Head', value=-1, minimum=-1, maximum=63, step=1, scale=1, interactive=True) # info='-1 to average on all heads'
+    return {
+        'position_block': __position,
+        'layer_block': __layer,
+        'head_block': __head,}
+
 # ACTIONS ######################################################################
 
 def create_actions_block() -> dict:
     __process = gradio.Button('Process', variant='primary', size='lg', scale=1, interactive=True)
-    __position = gradio.Slider(label='Position', value=-1, minimum=-1, maximum=128, step=1, scale=1, interactive=True) # info='-1 to average on all tokens'
-    return {
-        'process_block': __process,
-        'position_block': __position}
+    return {'process_block': __process,}
 
 # STATE ########################################################################
 
@@ -102,6 +106,8 @@ def create_layout(intro: str=INTRO) -> dict:
             with gradio.Row(equal_height=True):
                 __fields.update(create_inputs_block())
                 __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:
@@ -113,7 +119,7 @@ def create_layout(intro: str=INTRO) -> dict:
                     __fields.update(create_sampling_block())
                 with gradio.Row(equal_height=True):
                     __fields.update(create_target_block())
-                    __fields.update(create_display_block())
+                    # __fields.update(create_display_block())
     return __fields
 
 # EVENTS #######################################################################
@@ -121,57 +127,77 @@ def create_layout(intro: str=INTRO) -> dict:
 def update_layer_range(value: float, model: str) -> dict:
     return gradio.update(maximum=35, value=min(35, int(value))) if '120b' in model else gradio.update(maximum=23, value=min(23, int(value)))
 
-def update_position_range(value: float, tokens: list) -> dict:
-    return gradio.update(maximum=len(tokens) - 1, value=min(len(tokens) - 1, int(value)))
+def update_position_range(value: float, tokens: float) -> dict:
+    return gradio.update(maximum=int(tokens) - 1, value=min(int(tokens) - 1, int(value)))
 
 def update_computation_state(
     token_num: float,
     topk_num: float,
     topp_num: float,
+    token_idx: float,
+    layer_idx: float,
+    head_idx: float,
     prompt_str: str,
     device_str: str,
     model_obj: object,
     tokenizer_obj: object,
 ) -> tuple:
     # sanitize the inputs
-    __limit = max(1, min(128, int(token_num)))
-    __topk = max(1, min(128, int(token_num)))
-    __topp = max(0.0, min(1.0, float(token_num)))
-    __prompt = prompt_str.strip()
-    __device = device_str if (device_str in ['cpu', 'cuda']) else 'cpu'
+    __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)))
+    __layer_idx = max(-1, int(layer_idx))
+    __head_idx = max(-1, int(head_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))
     # handle all exceptions at once
     try:
         # dictionary {'input_ids': _, 'attention_mask': _}
-        __inputs = psaiops.score.attention.lib.preprocess_token_ids(
+        __input_data = psaiops.common.tokenizer.preprocess_token_ids(
             tokenizer_obj=tokenizer_obj,
-            prompt_str=__prompt,
-            device_str=__device)
+            prompt_str=__prompt_str,
+            device_str=__device_str)
         # parse the inputs
-        __input_dim = int(__inputs['input_ids'].shape[-1])
+        __input_dim = int(__input_data['input_ids'].shape[-1])
         # tensor (1, T)
-        __outputs = psaiops.score.attention.lib.generate_token_ids(
+        __output_data = psaiops.common.model.generate_token_ids(
             model_obj=model_obj,
-            input_args=__inputs,
-            token_num=__limit,
-            topk_num=__topk,
-            topp_num=__topp)
+            input_args=__input_data,
+            token_num=__token_num,
+            topk_num=__topk_num,
+            topp_num=__topp_num)
         # tensor (L, S, H, T, T)
-        __attentions = psaiops.score.attention.lib.compute_attention_weights(
+        __attention_data = psaiops.score.attention.lib.compute_attention_weights(
             model_obj=model_obj,
-            token_obj=__outputs)
+            token_obj=__output_data)
+        # reduce the layer, sample, head and output token axes => tensor (T,)
+        __score_data = psaiops.score.attention.lib.reduce_attention_weights(
+            attention_data=__attention_data,
+            token_idx=__token_idx,
+            layer_idx=__layer_idx,
+            head_idx=__head_idx,
+            input_dim=__input_dim)
+        # translate the scores into integer labels
+        __labels = psaiops.score.attention.lib.postprocess_attention_scores(
+            attention_data=__score_data,
+            input_dim=__input_dim,
+            token_idx=__token_idx)
         # detokenize the IDs
-        __tokens = psaiops.score.attention.lib.postprocess_token_ids(
+        __tokens = psaiops.common.tokenizer.postprocess_token_ids(
             tokenizer_obj=tokenizer_obj,
-            token_obj=__outputs)
-        # update each component => (input, output, attention) states
+            token_obj=__output_data)
+        # update each component => (input, output, attention, highligh) states
         return (
-            gradio.update(value=__tokens[:__input_dim]),
-            gradio.update(value=__tokens[__input_dim:]),
-            gradio.update(value=__attentions),)
+            list(zip(__tokens, __labels)),
+            __tokens[:__input_dim],
+            __tokens[__input_dim:],
+            __attention_data,)
     except:
         raise Exception('Attention generation aborted with an error.')
-    finally:
-        return (gradio.update(), gradio.update(), gradio.update())
 
 def update_text_highlight(
     token_idx: float,
@@ -180,15 +206,16 @@ def update_text_highlight(
     input_data: list,
     output_data: list,
     attention_data: torch.Tensor,
-) -> dict:
+) -> list:
     # sanitize the inputs
     __input_data = input_data or []
     __output_data = output_data or []
-    __attention_data = attention_data or torch.empty(0)
+    __attention_data = torch.empty(0) if (attention_data is None) else attention_data
     __input_dim = len(__input_data)
-    __token_idx = max(0, min(__input_dim, int(token_idx)))
-    __layer_idx = max(0, int(layer_idx))
-    __head_idx = max(0, int(head_idx))
+    __output_dim = len(__output_data)
+    __token_idx = max(-1, min(__output_dim, int(token_idx)))
+    __layer_idx = max(-1, int(layer_idx))
+    __head_idx = max(-1, int(head_idx))
     # exit if the data has not yet been computed
     if (not __input_data) or (not __output_data) or (attention_data is None) or (len(attention_data) == 0):
         return gradio.update()
@@ -209,11 +236,9 @@ def update_text_highlight(
             input_dim=__input_dim,
             token_idx=__token_idx)
         # update the component with [(token, label), ...]
-        return gradio.update(value=list(zip(__tokens, __labels)))
+        return list(zip(__tokens, __labels))
     except:
         raise Exception('Attention reduction aborted with an error.')
-    finally:
-        return gradio.update()
 
 # APP ##########################################################################
 
@@ -222,40 +247,49 @@ def create_app(title: str=TITLE, intro: str=INTRO, style: str=STYLE, model: str=
     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.score.attention.lib.get_model(name=model, device=__device)
-        __tokenizer = psaiops.score.attention.lib.get_tokenizer(name=model, device=__device)
+        __model = psaiops.common.model.get_model(name=model, device=__device)
+        __tokenizer = psaiops.common.tokenizer.get_tokenizer(name=model, device=__device)
         # adapt the computing function
         __compute = functools.partial(update_computation_state, model_obj=__model, tokenizer_obj=__tokenizer, device_str=__device)
         # create the UI
         __fields.update(create_layout(intro=intro))
         # init the state
         __fields.update(create_state())
-        # fetch the relevant fields
-        __button_block, __position_block, __output_block = (__fields['process_block'], __fields['position_block'], __fields['output_block'])
-        __output_state, __attention_state = (__fields['output_state'], __fields['attention_state'])
         # wire the input fields
-        __button_block.click(
+        __fields['tokens_block'].change(
+            fn=update_position_range,
+            inputs=[__fields[__k] for __k in ['position_block', 'tokens_block']],
+            outputs=__fields['position_block'],
+            queue=False,
+            show_progress='hidden')
+        __fields['model_block'].change(
+            fn=update_layer_range,
+            inputs=[__fields[__k] for __k in ['layer_block', 'model_block']],
+            outputs=__fields['layer_block'],
+            queue=False,
+            show_progress='hidden')
+        __fields['process_block'].click(
             fn=__compute,
-            inputs=[__fields[__k] for __k in ['tokens_block', 'topk_block', 'topp_block', 'input_block']],
-            outputs=[__fields[__k] for __k in ['input_state', 'output_state', 'attention_state']],
+            inputs=[__fields[__k] for __k in ['tokens_block', 'topk_block', 'topp_block', 'position_block', 'layer_block', 'head_block', 'input_block']],
+            outputs=[__fields[__k] for __k in ['output_block', 'input_state', 'output_state', 'attention_state']],
             queue=False,
             show_progress='full')
-        __output_state.change(
-            fn=update_position_range,
-            inputs=[__position_block, __output],
-            outputs=__position_block,
+        __fields['position_block'].change(
+            fn=update_text_highlight,
+            inputs=[__fields[__k] for __k in ['position_block', 'layer_block', 'head_block', 'input_state', 'output_state', 'attention_state']],
+            outputs=__fields['output_block'],
             queue=False,
             show_progress='hidden')
-        __attention_state.change(
+        __fields['layer_block'].change(
             fn=update_text_highlight,
             inputs=[__fields[__k] for __k in ['position_block', 'layer_block', 'head_block', 'input_state', 'output_state', 'attention_state']],
-            outputs=__output_block,
+            outputs=__fields['output_block'],
             queue=False,
             show_progress='hidden')
-        __position_block.change(
+        __fields['head_block'].change(
             fn=update_text_highlight,
             inputs=[__fields[__k] for __k in ['position_block', 'layer_block', 'head_block', 'input_state', 'output_state', 'attention_state']],
-            outputs=__output_block,
+            outputs=__fields['output_block'],
             queue=False,
             show_progress='hidden')
         # gradio application

psaiops/score/attention/lib.py

@@ -1,69 +1,7 @@
-import functools
-
 import torch
-import transformers
-
-import deformers.models.openai.gptoss
-
-# LOAD #########################################################################
-
-@functools.lru_cache(maxsize=4)
-def get_tokenizer(name: str, device: str='cpu'):
-    return transformers.AutoTokenizer.from_pretrained(
-        name,
-        use_fast=True,
-        dtype='auto',
-        device_map=device)
-
-@functools.lru_cache(maxsize=2)
-def get_model(name: str, device: str='cpu'):
-    __model = deformers.models.openai.gptoss.GptOssForCausalInference.from_pretrained(
-        name,
-        dtype='auto',
-        device_map=device)
-    # toggle the inference mode (not training)
-    __model.eval()
-    # transformers model
-    return __model
 
-# PREPROCESS #####################################################################
-
-@functools.lru_cache(maxsize=4)
-def preprocess_token_ids(
-    tokenizer_obj: object,
-    prompt_str: str,
-    device_str: str='cpu'
-) -> dict:
-    # tokenize
-    __inputs = tokenizer_obj(prompt_str, return_tensors='pt')
-    # move to the main device
-    return {__k: __v.to(device_str) for __k, __v in __inputs.items()}
-
-# GENERATE #######################################################################
-
-def generate_token_ids(
-    model_obj: object,
-    input_args: dict,
-    token_num: int,
-    topk_num: int = 4,
-    topp_num: float = 0.9,
-) -> torch.Tensor:
-    # 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=False,
-            output_attentions=False,
-            output_scores=False,
-            # early_stopping=True,
-            use_cache=True)
-    # full sequence
-    return __outputs.sequences # (1, T)
+import psaiops.common.model
+import psaiops.common.tokenizer
 
 # COMPUTE ########################################################################
 
@@ -91,14 +29,14 @@ def reduce_attention_weights(
 ) -> torch.Tensor:
     # parse
     __layer_dim, __batch_dim, __head_dim, __output_dim, __output_dim = tuple(attention_data.shape) # L, B, H, T, T
-    __layer_idx = min(layer_idx, __layer_dim)
-    __head_idx = min(head_idx, __head_dim)
+    __layer_idx = min(layer_idx, __layer_dim - 1)
+    __head_idx = min(head_idx, __head_dim - 1)
     __token_idx = min(token_idx, __output_dim - input_dim - 1) # T = I + O
     # select the relevant data along each axis
     __layer_slice = slice(None) if (__layer_idx < 0) else slice(__layer_idx, __layer_idx + 1)
     __sample_slice = slice(None)
     __head_slice = slice(None) if (__head_idx < 0) else slice(__head_idx, __head_idx + 1)
-    __token_slice = slice(input_dim, __output_dim) if (__token_idx < 0) else slice(input_dim + __token_idx, input_dim + __token_idx + 1)
+    __token_slice = slice(input_dim - 1, __output_dim) if (__token_idx < 0) else slice(input_dim + __token_idx - 1, input_dim + __token_idx)
     # filter the data
     __data = attention_data[__layer_slice, __sample_slice, __head_slice, __token_slice, slice(None)]
     # reduce all the axes but the last
@@ -119,7 +57,7 @@ def postprocess_attention_scores(
     __output_range = list(range(__output_dim - input_dim)) if (__token_idx < 0) else [__token_idx]
     __output_mask = torch.BoolTensor([__i in __output_range for __i in range(__output_dim - input_dim)])
     # normalize the scores
-    __input_scores = attention_data[__input_slice] / (attention_data[__input_slice].sum() + 1e-5)
+    __input_scores = attention_data[__input_slice] / (attention_data[__input_slice].max() + 1e-5)
     # round to obtain integer labels from 0 to 100
     __input_scores = torch.round(100.0 * __input_scores, decimals=0).type(torch.int32)
     # the generated tokens are not scored
@@ -127,19 +65,6 @@ def postprocess_attention_scores(
     # native list of serialized integers
     return [str(__i) for __i in __input_scores.tolist() + __output_scores.tolist()] # (I,) + (O,) = (T,)
 
-# POSTPROCESS ####################################################################
-
-def postprocess_token_ids(
-    tokenizer_obj: object,
-    token_obj: torch.Tensor,
-) -> list:
-    # remove the batch axis
-    __indices = token_obj.squeeze().tolist()
-    # back to token strings
-    __tokens = tokenizer_obj.convert_ids_to_tokens(__indices)
-    # normalize the tokens
-    return [__t.replace(chr(0x0120), ' ').replace(chr(0x010a), '\n') for __t in __tokens]
-
 # COMPUTE ########################################################################
 
 def score_tokens(
@@ -155,14 +80,14 @@ def score_tokens(
     tokenizer_obj: object,
 ) -> list:
     # dictionary {'input_ids': _, 'attention_mask': _}
-    __inputs = preprocess_token_ids(
+    __inputs = psaiops.common.tokenizer.preprocess_token_ids(
         tokenizer_obj=tokenizer_obj,
         prompt_str=prompt_str,
         device_str=device_str)
     # parse the inputs
     __input_dim = int(__inputs['input_ids'].shape[-1])
     # tensor (1, T)
-    __outputs = generate_token_ids(
+    __outputs = psaiops.common.tokenizer.model.generate_token_ids(
         model_obj=model_obj,
         input_args=__inputs,
         token_num=token_num,
@@ -185,7 +110,7 @@ def score_tokens(
         input_dim=__input_dim,
         token_idx=token_idx)
     # detokenize the IDs
-    __tokens = postprocess_token_ids(
+    __tokens = psaiops.common.tokenizer.postprocess_token_ids(
         tokenizer_obj=tokenizer_obj,
         token_obj=__outputs)
     # match tokens and labels for the HighlightedText field