psaiops 0.4.7__py3-none-any.whl

psaiops/score/router/app.py

@@ -0,0 +1,282 @@
+import functools
+
+import gradio
+import torch
+import torch.cuda
+import matplotlib.pyplot
+
+import psaiops.common.model
+import psaiops.common.tokenizer
+import psaiops.score.router.lib
+
+# META #########################################################################
+
+STYLE = '''.white-text span { color: white; }'''
+TITLE = '''Router 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 {
+        '0': '#000000',
+        '1': '#004444',}
+
+# 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}
+
+# PLOTS ########################################################################
+
+def create_plot_block() -> dict:
+    __plot = gradio.Plot(label='Router', scale=1)
+    return {'plot_block': __plot,}
+
+# OUTPUTS ######################################################################
+
+def create_outputs_block() -> dict:
+    __output = gradio.HighlightedText(label='Output', 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:
+    # __play = gradio.Button('>', variant='primary', size='lg', scale=1, interactive=True)
+    __position = gradio.Slider(label='Token', 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 {
+        'output_state': gradio.State(None),
+        'router_state': gradio.State(None),}
+
+# 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_plot_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(
+    current_val: float,
+    token_num: float,
+    output_data: torch.Tensor,
+) -> dict:
+    # take the generated tokens into account
+    __max = int(token_num) - 1 if (output_data is None) else int(output_data.shape[-1])
+    # keep the previous value if possible
+    __val = min(int(current_val), __max)
+    # return a gradio update dictionary
+    return gradio.update(maximum=__max, value=__val)
+
+def update_computation_state(
+    token_num: float,
+    topk_num: float,
+    topp_num: float,
+    token_idx: float,
+    prompt_str: str,
+    device_str: str,
+    model_obj: object,
+    tokenizer_obj: object,
+) -> tuple:
+    # sanitize the inputs
+    __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)))
+    __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), torch.empty(0))
+    # dictionary {'input_ids': _, 'attention_mask': _}
+    __input_data = psaiops.common.tokenizer.preprocess_token_ids(
+        tokenizer_obj=tokenizer_obj,
+        prompt_str=__prompt_str,
+        device_str=__device_str)
+    # tensor (1, T)
+    __output_data = psaiops.common.model.generate_token_ids(
+        model_obj=model_obj,
+        input_ids=__input_data['input_ids'],
+        attention_mask=__input_data['attention_mask'],
+        token_num=__token_num,
+        topk_num=__topk_num,
+        topp_num=__topp_num)
+    # tensor (L, S, H, T, T)
+    __router_data = psaiops.score.router.lib.compute_router_weights(
+        model_obj=model_obj,
+        token_data=__output_data)
+    # update each component => (highlight, plot) states
+    return (
+        __output_data.cpu(),
+        __router_data.cpu(),)
+
+def update_router_plot(
+    token_idx: float,
+    router_data: torch.Tensor,
+) -> tuple:
+    # exit if some values are missing
+    if (router_data is None) or (len(router_data) == 0):
+        return None
+    # reduce the batch and token axes => tensor (L, E)
+    __plot_data = psaiops.score.router.lib.reduce_router_weights(
+        router_data=router_data,
+        token_idx=int(token_idx),)
+    # translate the scores into integer labels
+    __plot_data = psaiops.score.router.lib.postprocess_router_weights(
+        router_data=__plot_data,)
+    # plot the data
+    __figure, __axes = matplotlib.pyplot.subplots()
+    __axes.imshow(__plot_data.float().numpy(), vmin=0.0, vmax=1.0, cmap='viridis')
+    __figure.tight_layout()
+    # remove the figure for the pyplot register for garbage collection
+    matplotlib.pyplot.close(__figure)
+    # update each component => (highlight, plot) states
+    return __figure
+
+def update_text_highlight(
+    token_idx: float,
+    output_data: torch.Tensor,
+    tokenizer_obj: object,
+) -> list:
+    # exit if some values are missing
+    if (output_data is None) or (len(output_data) == 0):
+        return None
+    # detokenize the IDs
+    __token_str = psaiops.common.tokenizer.postprocess_token_ids(
+        tokenizer_obj=tokenizer_obj,
+        token_data=output_data)
+    # list of string classes
+    __token_cls = psaiops.score.router.lib.postprocess_token_cls(
+        token_idx=int(token_idx),
+        token_dim=len(__token_str))
+    # pairs of token and class
+    return list(zip(__token_str, __token_cls))
+
+# 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.common.model.get_model(name=model, device=__device)
+        __tokenizer = psaiops.common.tokenizer.get_tokenizer(name=model, device=__device)
+        # adapt the event handlers
+        __compute = functools.partial(update_computation_state, model_obj=__model, tokenizer_obj=__tokenizer, device_str=__device)
+        __highlight = functools.partial(update_text_highlight, tokenizer_obj=__tokenizer)
+        # create the UI
+        __fields.update(create_layout(intro=intro))
+        # init the state
+        __fields.update(create_state())
+        # update the data after clicking process
+        __fields['process_block'].click(
+            fn=__compute,
+            inputs=[__fields[__k] for __k in ['tokens_block', 'topk_block', 'topp_block', 'position_block', 'input_block']],
+            outputs=[__fields[__k] for __k in ['output_state', 'router_state']],
+            queue=False,
+            show_progress='full').then(
+        # update the range of the position slider when the output changes
+            fn=update_position_range,
+            inputs=[__fields[__k] for __k in ['position_block', 'tokens_block', 'output_state']],
+            outputs=__fields['position_block'],
+            queue=False,
+            show_progress='hidden').then(
+        # update the token highlight when the output data changes
+            fn=__highlight,
+            inputs=[__fields[__k] for __k in ['position_block', 'output_state']],
+            outputs=__fields['output_block'],
+            queue=False,
+            show_progress='full').then(
+        # update the plot when the router data changes
+            fn=update_router_plot,
+            inputs=[__fields[__k] for __k in ['position_block', 'router_state']],
+            outputs=__fields['plot_block'],
+            queue=False,
+            show_progress='full')
+        # update the range of the position slider when the settings change
+        __fields['tokens_block'].change(
+            fn=update_position_range,
+            inputs=[__fields[__k] for __k in ['position_block', 'tokens_block', 'output_state']],
+            outputs=__fields['position_block'],
+            queue=False,
+            show_progress='hidden')
+        # update the plot when the focus changes
+        __fields['position_block'].change(
+            fn=update_router_plot,
+            inputs=[__fields[__k] for __k in ['position_block', 'router_state']],
+            outputs=__fields['plot_block'],
+            queue=False,
+            show_progress='full')
+        # update the token highlight when the token focus changes
+        __fields['position_block'].change(
+            fn=__highlight,
+            inputs=[__fields[__k] for __k in ['position_block', 'output_state']],
+            outputs=__fields['output_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,59 @@
+import functools
+
+import torch
+
+# COMPUTE ########################################################################
+
+def compute_router_weights(
+    model_obj: object,
+    token_data: torch.Tensor,
+) -> torch.Tensor:
+    # process the full sequence
+    with torch.no_grad():
+        __outputs = model_obj(
+            input_ids=token_data,
+            output_attentions=False,
+            output_router_logits=True,
+            return_dict=True)
+    # stack all the layer outputs L * (T, E) => (L, T, E)
+    __logits = torch.stack(__outputs.router_logits, dim=0)
+    # turn the logits into expert probabilities
+    return torch.softmax(__logits, dim=-1)
+
+# REDUCE #######################################################################
+
+def reduce_router_weights(
+    router_data: torch.Tensor,
+    token_idx: int, # -1 => avg over all tokens
+) -> torch.Tensor:
+    # parse
+    __layer_dim, __token_dim, __expert_dim = tuple(router_data.shape) # L, T, E
+    __token_idx = min(token_idx, __token_dim - 1)
+    # select the relevant data along each axis
+    __token_slice = slice(0, __token_dim) if (__token_idx < 0) else slice(__token_idx, __token_idx + 1)
+    # filter the data
+    __data = router_data[slice(None), __token_slice, slice(None)]
+    # reduce all the axes but the last
+    return __data.mean(dim=1, keepdim=False)
+
+# FORMAT #########################################################################
+
+def postprocess_router_weights(
+    router_data: torch.Tensor, # (L, E)
+) -> list:
+    # the averaging over tokens may have broken the scaling
+    __probs = torch.softmax(router_data, dim=-1)
+    # enforce the output range [0; 1] with 1 included
+    return __probs / __probs.amax(dim=-1, keepdim=True)
+
+# POSTPROCESS ####################################################################
+
+def postprocess_token_cls(
+    token_idx: int,
+    token_dim: int,
+) -> list:
+    __token_idx = max(-1, min(token_dim, token_idx))
+    # class 1 for the focused token(s) 0 for the rest
+    __token_cls = [str(int(__i == token_idx)) for __i in range(token_dim)]
+    # average on all the tokens when the idx is negative
+    return token_dim * ['1'] if (token_idx < 0) else __token_cls

psaiops/score/shapley/app.py

@@ -0,0 +1,158 @@
+import functools
+
+import gradio
+import torch
+import torch.cuda
+
+import psaiops.common.model
+import psaiops.common.tokenizer
+
+# META #########################################################################
+
+STYLE = '''.white-text span { color: white; }'''
+TITLE = '''Shapley 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.common.model.get_model(name=model, device=__device)
+        __tokenizer = psaiops.common.tokenizer.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 @@
+import torch

psaiops/score/similarity/app.py

@@ -0,0 +1,152 @@
+import functools
+
+import gradio
+import torch
+import torch.cuda
+
+import psaiops.common.model
+import psaiops.common.tokenizer
+
+# META #########################################################################
+
+STYLE = '''.white-text span { color: white; }'''
+TITLE = '''Similarity Scoring'''
+INTRO = '''Score each token according to its similarity with a reference 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:
+    __reference = gradio.Textbox(label='Reference', value='', placeholder='The text used as reference for the comparison.', lines=4, scale=1, show_copy_button=True, interactive=True)
+    __prompt = gradio.Textbox(label='Prompt', value='', placeholder='A string of tokens to compare to the reference.', lines=4, scale=1, show_copy_button=True, interactive=True)
+    return {
+        'reference_input_block': __reference,
+        'prompt_input_block': __prompt,}
+
+# OUTPUTS ######################################################################
+
+def create_outputs_block() -> dict:
+    __reference = gradio.HighlightedText(label='', value='', scale=1, interactive=False, show_label=False, show_legend=False, show_inline_category=False, combine_adjacent=False, color_map=create_color_map(), elem_classes='white-text')
+    __prompt = gradio.HighlightedText(label='', value='', scale=1, interactive=False, show_label=False, show_legend=False, show_inline_category=False, combine_adjacent=False, color_map=create_color_map(), elem_classes='white-text')
+    return {
+        'reference_output_block': __reference,
+        'prompt_output_block': __prompt,}
+
+# 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())
+            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_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.common.model.get_model(name=model, device=__device)
+        __tokenizer = psaiops.common.tokenizer.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/similarity/lib.py

@@ -0,0 +1 @@
+import torch