psaiops 0.1.1__tar.gz → 0.2.0__tar.gz

{psaiops-0.1.1 → psaiops-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: psaiops
-Version: 0.1.1
+Version: 0.2.0
 Summary: Web apps to inspect & engineer NN activations.
 License: .github/LICENSE.md
 Author: apehex
@@ -12,9 +12,12 @@ Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13
+Requires-Dist: accelerate (>=1.10)
 Requires-Dist: deformers (>=0.0)
 Requires-Dist: gradio (>=5.0)
+Requires-Dist: kernels (==0.10)
 Requires-Dist: requests (>=2.0)
+Requires-Dist: triton (==3.4)
 Description-Content-Type: text/markdown
 
 # psAI ops <img src="images/logo.png" alt="apehex logo" width="32" height="32">

psaiops-0.2.0/psaiops/common/dropdown.py

@@ -0,0 +1,19 @@
+import gradio
+
+import psaiops.elements.data
+
+# AUTO-COMPLETE ################################################################
+
+def update_dropdown(label: str, data: gradio.KeyUpData):
+    # model_dropdown.key_up(fn=update_dropdown, inputs=[model_dropdown, gradio.State("model")], outputs=model_dropdown, queue=False, show_progress="hidden")
+    datasets = psaiops.elements.data.query_huggingface(target=data.input_value, label=label, limit=16)
+    return gradio.update(choices=datasets, visible=True)
+
+# with gradio.Blocks() as demo:
+#     model_dropdown = gradio.Dropdown(label="Models Auto-Complete", choices=[""], allow_custom_value=True)
+#     dataset_dropdown = gradio.Dropdown(label="Datasets Auto-Complete", choices=[""], allow_custom_value=True)
+#     spaces_dropdown = gradio.Dropdown(label="Spaces Auto-Complete", choices=[""], allow_custom_value=True)
+#     model_dropdown.key_up(fn=update_dropdown, inputs=[gradio.State("model")], outputs=model_dropdown, queue=False, show_progress="hidden")
+#     dataset_dropdown.key_up(fn=update_dropdown, inputs=[gradio.State("dataset")], outputs=dataset_dropdown, queue=False, show_progress="hidden")
+#     spaces_dropdown.key_up(fn=update_dropdown, inputs=[gradio.State("space")], outputs=spaces_dropdown, queue=False, show_progress="hidden")
+# demo.launch(share=True, debug=True)

psaiops-0.2.0/psaiops/compose/contrast/app.py

@@ -0,0 +1,178 @@
+import functools
+
+import gradio
+import torch
+import torch.cuda
+
+import psaiops.compose.contrast.lib
+
+# META #########################################################################
+
+TITLE = '''Contrastive Steering'''
+INTRO = '''Add a delta of activation to a prompt to steer the model output in a specific latent direction.'''
+STYLE = '''.giga-text input { font-size: 32px; }'''
+
+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)
+    return {'intro_block': __intro}
+
+# MODEL ########################################################################
+
+def create_model_block() -> dict:
+    __model = gradio.Dropdown(label='Model ID', 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 = gradio.Slider(label='Layer Depth', value=12, minimum=0, maximum=23, step=1, scale=1, interactive=True)
+    return {
+        'model_block': __model,
+        'layer_block': __layer,}
+
+# 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,}
+
+# REDUCTION ####################################################################
+
+def create_reduction_block() -> dict:
+    __from = gradio.Slider(label='Average From', value=0, minimum=0, maximum=256, step=1, scale=1, interactive=True)
+    __to = gradio.Slider(label='Average To', value=256, minimum=0, maximum=256, step=1, scale=1, interactive=True)
+    return {
+        'from_block': __from,
+        'to_block': __to,}
+
+# INPUTS #######################################################################
+
+def create_inputs_row(operation: str='', index: int=0, label: bool=False) -> dict:
+    # __operation = gradio.Button(value=operation, variant='primary', size='lg', elem_classes='white-text', scale=1, interactive=False)
+    __operation = gradio.Dropdown(label=f'Operation', value=operation, choices=['', '+', '-', 'x', '.', '='], elem_classes='giga-text', scale=1, show_label=label, allow_custom_value=False, multiselect=False, interactive=False)
+    __alpha = gradio.Slider(label='Factor', value=1.0, minimum=0.0, maximum=16.0, step=0.1, scale=1, show_label=label, interactive=True)
+    __input = gradio.Textbox(label=f'Prompt', value='', placeholder='Some text.', lines=2, max_lines=2, scale=8, show_label=label, show_copy_button=True, interactive=True)
+    return {
+        f'operation_{index}_block': __operation,
+        f'factor_{index}_block': __alpha,
+        f'prompt_{index}_block': __input,}
+
+# OUTPUTS ######################################################################
+
+def create_outputs_block() -> dict:
+    __output = gradio.Textbox(label='= Total', value='', placeholder='Some text.', lines=2, max_lines=8, scale=1, show_label=True, show_copy_button=True, interactive=False)
+    return {'output_block': __output}
+
+# ACTIONS ######################################################################
+
+def create_actions_block() -> dict:
+    __process = gradio.Button(value='Process', variant='primary', size='lg', scale=1, interactive=True)
+    return {'process_block': __process,}
+
+# TABLE ########################################################################
+
+def create_table_block() -> dict:
+    __table = gradio.DataFrame(label='Summary', type='numpy', headers=None,  row_count=4, col_count=256, scale=1, interactive=False)
+    return {'table_block': __table,}
+
+# 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('Equation') as __main_tab:
+            __fields.update({'main_tab': __main_tab})
+            with gradio.Row(equal_height=True):
+                __fields.update(create_inputs_row(operation='', index=0, label=True))
+            with gradio.Row(equal_height=True):
+                __fields.update(create_inputs_row(operation='-', index=1, label=False))
+            with gradio.Row(equal_height=True):
+                __fields.update(create_inputs_row(operation='+', index=2, label=False))
+            with gradio.Row(equal_height=True):
+                __fields.update(create_outputs_block())
+            with gradio.Row(equal_height=True):
+                __fields.update(create_actions_block())
+        with gradio.Tab('Details') as __details_tab:
+            with gradio.Row(equal_height=True):
+                __fields.update(create_table_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_reduction_block())
+                    # __fields.update(create_display_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_table_data(positive: str, negative: str, prompt: str, output: str, tokenizer: object) -> list:
+    __outputs = tokenizer([positive, negative, prompt, output], padding=True)
+    return [tokenizer.convert_ids_to_tokens(__s) for __s in __outputs['input_ids']]
+
+# 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.compose.contrast.lib.get_model(name=model, device=__device)
+        __tokenizer = psaiops.compose.contrast.lib.get_tokenizer(name=model, device=__device)
+        # adapt the computing functions
+        __compute = functools.partial(psaiops.compose.contrast.lib.steer_model_output, model_obj=__model, tokenizer_obj=__tokenizer, device_str=__device)
+        __format = functools.partial(update_table_data, tokenizer=__tokenizer)
+        # create the UI
+        __fields.update(create_layout(intro=intro))
+        # init the state
+        __fields.update(create_state())
+        # wire the input fields
+        __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['output_block'].change(
+            fn=__format,
+            inputs=[__fields[__k] for __k in ['prompt_0_block', 'prompt_1_block', 'prompt_2_block', 'output_block']],
+            outputs=__fields['table_block'],
+            queue=False,
+            show_progress='hidden')
+        __fields['process_block'].click(
+            fn=__compute,
+            inputs=[__fields[__k] for __k in ['prompt_0_block', 'prompt_1_block', 'prompt_2_block', 'factor_0_block', 'factor_1_block', 'factor_2_block', 'tokens_block', 'topk_block', 'topp_block', 'layer_block']],
+            outputs=__fields['output_block'],
+            queue=False,
+            show_progress='full')
+        # gradio application
+        return __app
+
+# MAIN #########################################################################
+
+if __name__ == '__main__':
+    __app = create_app()
+    __app.launch(share=True, debug=True)

psaiops-0.2.0/psaiops/compose/contrast/lib.py

@@ -0,0 +1,175 @@
+import functools
+
+import torch
+import torch.nn
+import torch.nn.modules
+import transformers
+
+import deformers.models.openai.gptoss
+import mlable.shapes
+
+# 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: object,
+    prompts: list,
+    device: str='cpu'
+) -> dict:
+    # tokenize
+    __inputs = tokenizer(prompts, return_tensors='pt', padding=True)
+    # move to the main device
+    return {__k: __v.to(device) for __k, __v in __inputs.items()}
+
+# HOOK #########################################################################
+
+def capture_hidden_activation(
+    module: torch.nn.modules.Module,
+    inputs: torch.Tensor,
+    outputs: torch.Tensor,
+    index: int,
+    captured: dict,
+) -> None:
+    captured[index] = outputs # (B, S, E)
+
+# MASKS ########################################################################
+
+def compute_sequence_mask(
+    tokens: torch.Tensor, # (B, S)
+    masks: torch.Tensor, # (B, S)
+) -> torch.Tensor:
+    # keep the intersection of the masks
+    __masks = masks.prod(dim=0, keepdim=False)
+    # and only the positions that differ between positive and negative prompts
+    return __masks * (tokens[0] != tokens[1])
+
+# REDUCTION ####################################################################
+
+def compute_delta_activation(
+    data: torch.Tensor, # (B, S, E)
+    masks: torch.Tensor, # (S,)
+    signs: torch.Tensor, # (B,)
+    keepdim: bool=True,
+) -> torch.Tensor:
+    __dtype = data.dtype
+    __device = data.device
+    # sign each sample along the batch axis
+    __shape = tuple(mlable.shapes.filter(data.shape, axes=[0]))
+    __signs = signs.to(dtype=__dtype, device=__device).view(__shape)
+    # combine along the batch axis to keep the shortest mask on the sequence axis
+    __shape = tuple(mlable.shapes.filter(data.shape, axes=[1]))
+    __masks = masks.to(dtype=__dtype, device=__device).view(__shape)
+    # mean factor: half the signs size along the batch axis and the number of positions kept along the sequence axis
+    __factor = (0.5 * float(len(__signs)) * __masks.sum()).clamp(min=1e-8)
+    # take the difference along the batch axis and the average along the sequence axis
+    return (data * __signs * __masks).sum(dim=[0, 1], keepdim=keepdim) / __factor
+
+# DELTA ########################################################################
+
+def add_delta_activation(
+    module: torch.nn.modules.Module,
+    inputs: torch.Tensor,
+    outputs: torch.Tensor,
+    delta: torch.Tensor,
+    alpha: torch.Tensor,
+    beta: torch.Tensor,
+) -> torch.Tensor:
+    # expand the single feature axis of the delta
+    __shape = mlable.shapes.filter(outputs.shape, axes=[-1])
+    # rescale the delta
+    return alpha * outputs + beta * delta.view(__shape)
+
+# MAIN #########################################################################
+
+def steer_model_output(
+    positive_str: str,
+    negative_str: str,
+    prompt_str: str,
+    positive_rate: float,
+    negative_rate: float,
+    prompt_rate: float,
+    token_num: int,
+    topk_num: int,
+    topp_num: float,
+    layer_idx: int,
+    device_str: str,
+    model_obj: object,
+    tokenizer_obj: object,
+) -> str:
+    # parse & sanitize
+    __prompt0 = positive_str.strip()
+    __prompt1 = negative_str.strip()
+    __prompt2 = prompt_str.strip()
+    __alpha0 = max(0.0, float(positive_rate))
+    __alpha1 = max(0.0, float(negative_rate))
+    __alpha2 = max(0.0, float(prompt_rate))
+    __count = max(1, int(token_num))
+    __topk = max(1, int(topk_num))
+    __topp = max(0.0, float(topp_num))
+    __index = max(0, int(layer_idx))
+    # store hidden states
+    __captured = {}
+    # stop if inputs are missing
+    if not (__prompt0 and __prompt1 and __prompt2):
+        return ''
+    # tokenize the 2 prompts and pad to same length
+    __inputs = preprocess_token_ids(tokenizer=tokenizer_obj, prompts=(__prompt0, __prompt1), device=device_str)
+    # forward hook to capture output hidden state
+    __hook = functools.partial(capture_hidden_activation, index=__index, captured=__captured)
+    # attach to the model
+    __handle = model_obj.model.layers[__index].register_forward_hook(__hook)
+    with torch.no_grad():
+        # inference mode
+        model_obj.eval().to(device_str)
+        # prefill with a single forward
+        __outputs = model_obj(**__inputs, use_cache=True, output_attentions=False, output_hidden_states=False, return_dict=True)
+    # stop capturing activations
+    __handle.remove()
+    # select only the positions where the tokens differ
+    __masks = compute_sequence_mask(tokens=__inputs['input_ids'], masks=__inputs['attention_mask'])
+    # activation delta at layer L
+    __delta = compute_delta_activation(data=__captured[__index], masks=__masks, signs=torch.Tensor([1, -1]), keepdim=False)
+    # add the delta on every forward pass
+    __hook = functools.partial(add_delta_activation, alpha=__alpha2, beta=0.5 * (__alpha0 + __alpha1), delta=__delta)
+    # attach to the model
+    __handle = model_obj.model.layers[__index].register_forward_hook(__hook)
+    # now process the user input
+    __inputs = preprocess_token_ids(tokenizer=tokenizer_obj, prompts=(prompt_str,), device=device_str)
+    # generate the new with tampered activations
+    with torch.no_grad():
+        __outputs = model_obj.generate(
+            **__inputs,
+            max_new_tokens=__count,
+            do_sample=(0.0 < __topp < 1.0) or (__topk > 0),
+            top_k=__topk if (__topk > 0) else None,
+            top_p=__topp if (0.0 < __topp <= 1.0) else None,
+            return_dict_in_generate=True,
+            output_hidden_states=False,
+            output_attentions=False,
+            output_scores=False,
+            use_cache=True)
+    # stop altering the activations
+    __handle.remove()
+    # single string
+    return tokenizer_obj.decode(__outputs.sequences[0], skip_special_tokens=True)

{psaiops-0.1.1 → psaiops-0.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "psaiops"
-version = "0.1.1"
+version = "0.2.0"
 description = "Web apps to inspect & engineer NN activations."
 license = ".github/LICENSE.md"
 readme = ".github/README.md"
@@ -9,9 +9,12 @@ packages = [{include = "psaiops"}]
 
 [tool.poetry.dependencies]
 python = ">=3.10, <3.14"
+accelerate = ">=1.10"
 deformers = ">=0.0"
 gradio = ">=5.0"
+kernels = "0.10"
 requests = ">=2.0"
+triton = "==3.4"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "*"