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

psaiops/score/attention/app.py

@@ -128,50 +128,69 @@ 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(0, min(__token_num, int(token_idx)))
+    __layer_idx = max(0, int(layer_idx))
+    __head_idx = max(0, int(head_idx))
+    __prompt_str = prompt_str.strip()
+    __device_str = device_str if (device_str in ['cpu', 'cuda']) else 'cpu'
     # handle all exceptions at once
     try:
         # dictionary {'input_ids': _, 'attention_mask': _}
-        __inputs = psaiops.score.attention.lib.preprocess_token_ids(
+        __input_data = psaiops.score.attention.lib.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.score.attention.lib.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(
             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),)
+            gradio.update(value=__attention_data),
+            gradio.update(value=list(zip(__tokens, __labels))))
     except:
         raise Exception('Attention generation aborted with an error.')
     finally:
-        return (gradio.update(), gradio.update(), gradio.update())
+        return (gradio.update(), gradio.update(), gradio.update(), gradio.update())
 
 def update_text_highlight(
     token_idx: float,
@@ -236,22 +255,16 @@ def create_app(title: str=TITLE, intro: str=INTRO, style: str=STYLE, model: str=
         # wire the input fields
         __button_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 ['input_state', 'output_state', 'attention_state', 'output_block']],
             queue=False,
             show_progress='full')
         __output_state.change(
             fn=update_position_range,
-            inputs=[__position_block, __output],
+            inputs=[__position_block, __output_state],
             outputs=__position_block,
             queue=False,
             show_progress='hidden')
-        __attention_state.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,
-            queue=False,
-            show_progress='hidden')
         __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']],

{psaiops-0.0.13.dist-info → psaiops-0.0.15.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: psaiops
-Version: 0.0.13
+Version: 0.0.15
 Summary: Web apps to inspect & engineer NN activations.
 License: .github/LICENSE.md
 Author: apehex

{psaiops-0.0.13.dist-info → psaiops-0.0.15.dist-info}/RECORD

@@ -7,9 +7,9 @@ psaiops/elements/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,
 psaiops/elements/data.py,sha256=vGYeMN11uP9gs8rV6aSDffE_TeIX5PmdzWGwUpdGE2Y,906
 psaiops/score/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 psaiops/score/attention/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-psaiops/score/attention/app.py,sha256=NB0LVbRb-B9HniSHRJSg3Z893DxhGh06UWwN1A1KKv4,11961
+psaiops/score/attention/app.py,sha256=ylPOko30zqcJddNbqc2z3ZrUUqNxjw1EFbdZJbODUJ8,12655
 psaiops/score/attention/lib.py,sha256=kPLdmWCYbpMnd7fi4ZeCvxiGzJuvTKYSXgRAm_w02KA,6928
 psaiops/steer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-psaiops-0.0.13.dist-info/METADATA,sha256=kU4BWy4WEUlWemYWIZ8lDP1IQp8_T8twMP-cy5a8SV8,1222
-psaiops-0.0.13.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-psaiops-0.0.13.dist-info/RECORD,,
+psaiops-0.0.15.dist-info/METADATA,sha256=Ex3PrfxidXsUR1LOEi_zXpOj1nhBcCIDUAY-RkQsvWA,1222
+psaiops-0.0.15.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+psaiops-0.0.15.dist-info/RECORD,,