psaiops 0.2.0__tar.gz → 0.2.1__tar.gz

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

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

{psaiops-0.2.0 → psaiops-0.2.1}/psaiops/compose/contrast/app.py

@@ -1,6 +1,7 @@
 import functools
 
 import gradio
+import pandas
 import torch
 import torch.cuda
 
@@ -129,9 +130,22 @@ 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_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']]
+def update_table_data(positive: str, negative: str, prompt: str, output: str, tokenizer: object) -> pandas.DataFrame:
+    # array of token IDs
+    __outputs = tokenizer([positive, negative, prompt, output], return_tensors='pt', padding=True)
+    # array of token strings
+    __tokens = [tokenizer.convert_ids_to_tokens(__s) for __s in __outputs['input_ids']]
+    # shift the special characters
+    __tokens = [[__t.replace(chr(0x0120), ' ').replace(chr(0x010a), '\\n') for __t in __s] for __s in __tokens]
+    # mask the tokens that differ between positive and negative prompts
+    __masks = psaiops.compose.contrast.lib.compute_sequence_mask(tokens=__outputs['input_ids'])
+    # convert into a data frame
+    __data = pandas.DataFrame(__tokens)
+    # color the background in red for the positions marked by the mask
+    return __data.style.apply(update_table_style, masks=pandas.DataFrame(__masks), axis=None)
+
+def update_table_style(data: pandas.DataFrame, masks: pandas.DataFrame) -> pandas.DataFrame:
+    return pandas.DataFrame(masks.replace({True: 'background-color: rgb(255, 0, 0, 64%)', False: 'background-color: rgb(0, 0, 0, 0%)',}))
 
 # APP ##########################################################################
 

{psaiops-0.2.0 → psaiops-0.2.1}/psaiops/compose/contrast/lib.py

@@ -57,33 +57,37 @@ def capture_hidden_activation(
 
 def compute_sequence_mask(
     tokens: torch.Tensor, # (B, S)
-    masks: 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])
+    __shape = mlable.shapes.divide(tokens.shape, axis=0, factor=2, insert=True)
+    # group the samples two by two
+    __data = tokens.reshape(__shape)
+    # compare each sample with its neighbor
+    __masks = __data[:, :1]  != __data[:, 1:]
+    # apply the same mask to both samples
+    return __masks.expand(__shape).reshape(tokens.shape)
 
 # REDUCTION ####################################################################
 
 def compute_delta_activation(
     data: torch.Tensor, # (B, S, E)
-    masks: torch.Tensor, # (S,)
+    masks: torch.Tensor, # (B, S,)
     signs: torch.Tensor, # (B,)
     keepdim: bool=True,
 ) -> torch.Tensor:
     __dtype = data.dtype
     __device = data.device
+    __dim0, __dim1, __dim2 = tuple(data.shape)
     # 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]))
+    __shape = tuple(mlable.shapes.filter(data.shape, axes=[0, 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)
+    __factor = (0.5 * float(__dim0) * __masks.sum(dim=1, keepdim=True)).clamp(min=1.0)
     # 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
+    return (data * __signs * __masks / __factor).sum(dim=[0, 1], keepdim=keepdim)
 
 # DELTA ########################################################################
 
@@ -147,7 +151,7 @@ def steer_model_output(
     # 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'])
+    __masks = compute_sequence_mask(tokens=__inputs['input_ids'])
     # 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

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

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "psaiops"
-version = "0.2.0"
+version = "0.2.1"
 description = "Web apps to inspect & engineer NN activations."
 license = ".github/LICENSE.md"
 readme = ".github/README.md"