psaiops 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

psaiops/score/router/app.py

@@ -0,0 +1,133 @@
+import functools
+
+import gradio
+import torch
+import torch.cuda
+
+import psaiops.score.router.lib
+
+# META #########################################################################
+
+STYLE = '''.white-text span { color: white; }'''
+TITLE = '''Attention 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 {
+        '-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, 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}
+
+# OUTPUTS ######################################################################
+
+def create_outputs_block() -> dict:
+    __output = gradio.Plot(label='Logits', scale=1)
+    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'
+    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 {}
+
+# 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_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(value: float, tokens: float) -> dict:
+    return gradio.update(maximum=int(tokens) - 1, value=min(int(tokens) - 1, int(value)))
+
+# 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.score.router.lib.get_model(name=model, device=__device)
+        __tokenizer = psaiops.score.router.lib.get_tokenizer(name=model, device=__device)
+        # create the UI
+        __fields.update(create_layout(intro=intro))
+        # init the state
+        __fields.update(create_state())
+        # wire the input fields
+        __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')
+        # gradio application
+        return __app
+
+# MAIN #########################################################################
+
+if __name__ == '__main__':
+    __app = create_app()
+    __app.launch(share=True, debug=True)

psaiops/score/router/lib.py

@@ -0,0 +1,66 @@
+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)

psaiops/score/shapley/app.py

@@ -0,0 +1,157 @@
+import functools
+
+import gradio
+import torch
+import torch.cuda
+
+import psaiops.score.shapley.lib
+
+# META #########################################################################
+
+STYLE = '''.white-text span { color: white; }'''
+TITLE = '''Attention Scoring'''
+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.'''
+
+MODEL = 'openai/gpt-oss-20b'
+
+# COLORS #######################################################################
+
+def create_color_map() -> dict:
+    return {
+        '-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, 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,}
+
+# SAMPLING #####################################################################
+
+def create_computation_block() -> dict:
+    __count = gradio.Slider(label='Sample count', value=4, minimum=1, maximum=16, step=1, scale=1, interactive=True)
+    __min = gradio.Slider(label='Min size', value=50, minimum=0, maximum=100, step=1, scale=1, interactive=True)
+    __max = gradio.Slider(label='Max size', value=100, minimum=0, maximum=100, step=1, scale=1, interactive=True)
+    return {
+        'count_block': __count,
+        'min_block': __min,
+        'max_block': __max,}
+
+# 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}
+
+# 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=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'
+    return {
+        'position_block': __position,
+        'layer_block': __layer,}
+
+# 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 {}
+
+# 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())
+                __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())
+                with gradio.Row(equal_height=True):
+                    __fields.update(create_computation_block())
+    return __fields
+
+# EVENTS #######################################################################
+
+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: float) -> dict:
+    return gradio.update(maximum=int(tokens) - 1, value=min(int(tokens) - 1, int(value)))
+
+# 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.score.shapley.lib.get_model(name=model, device=__device)
+        __tokenizer = psaiops.score.shapley.lib.get_tokenizer(name=model, device=__device)
+        # create the UI
+        __fields.update(create_layout(intro=intro))
+        # init the state
+        __fields.update(create_state())
+        # wire the input fields
+        __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')
+        # gradio application
+        return __app
+
+# MAIN #########################################################################
+
+if __name__ == '__main__':
+    __app = create_app()
+    __app.launch(share=True, debug=True)

psaiops/score/shapley/lib.py

@@ -0,0 +1,66 @@
+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)

{psaiops-0.3.0.dist-info → psaiops-0.3.2.dist-info}/METADATA

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

{psaiops-0.3.0.dist-info → psaiops-0.3.2.dist-info}/RECORD

@@ -16,7 +16,13 @@ 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=UK6-11xUP9S7_8ryZiqFaGfJvG_scS-NzU6QJLEVbEs,13479
 psaiops/score/attention/lib.py,sha256=j9p7ZXhkh5Y9VDeIe6TPUUrHRkwp1Y6P22iciaaKdL4,6932
+psaiops/score/router/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+psaiops/score/router/app.py,sha256=v2z4Qjho4naVmPandsX2nWwaCElwaWHnSVM0g5UDG1c,5640
+psaiops/score/router/lib.py,sha256=Cra3jksyqCunCRLsQ88ts0RITU7kC8LgudjIdWfhmyE,1996
+psaiops/score/shapley/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+psaiops/score/shapley/app.py,sha256=eZX-Hw_tj9x7lXyazxVMwBlRPaOHS6jxWcnOdnLzAz4,7151
+psaiops/score/shapley/lib.py,sha256=Cra3jksyqCunCRLsQ88ts0RITU7kC8LgudjIdWfhmyE,1996
 psaiops/steer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-psaiops-0.3.0.dist-info/METADATA,sha256=lKbJIDRYFfShvGp8gSFuYtqKKhZW_e6tbcMXdgaodSs,1318
-psaiops-0.3.0.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
-psaiops-0.3.0.dist-info/RECORD,,
+psaiops-0.3.2.dist-info/METADATA,sha256=jM6uT4ydy8THc_LZftgwj23aruOxmh0QrJqXsWclNtM,1318
+psaiops-0.3.2.dist-info/WHEEL,sha256=fGIA9gx4Qxk2KDKeNJCbOEwSrmLtjWCwzBz351GyrPQ,88
+psaiops-0.3.2.dist-info/RECORD,,