Create app.py

app.py

@@ -0,0 +1,432 @@
+import csv
+import io
+from pathlib import Path
+
+import numpy as np
+import plotly.graph_objects as go
+import streamlit as st
+from plotly.colors import qualitative
+
+from embeddings import Model, embedding_name, load_tokenizer, nearest_tokens, search_embeddings
+from relationships import relationship_figure
+from relationship_ui import render_relationship_controls
+from model_picker import show_model_picker
+from hub_ui import render_hub_loader
+from upload_ui import render_upload_loader
+from app_environment import hosted_mode, remove_session_directory, UPLOAD_ROOT, PRIVATE_CACHE_ROOT
+
+
+st.set_page_config(page_title='Token Visualiser', page_icon='🔬', layout='wide', initial_sidebar_state='collapsed')
+hosted = hosted_mode()
+if st.session_state.pop('clear_hub_token', False):
+    st.session_state.pop('hf_token', None)
+if 'pending_chart_view' in st.session_state:
+    st.session_state.chart_view = st.session_state.pop('pending_chart_view')
+if 'pending_model' in st.session_state:
+    pending_model = st.session_state.pop('pending_model')
+    if not hosted:
+        st.session_state.model_location = pending_model
+
+
+def clear_model():
+    # Keep the last browsing location, but release model data and all analysis state.
+    model = st.session_state.get('model')
+    if model is not None:
+        model.close()
+    for key, root in [('upload_directory', UPLOAD_ROOT), ('private_hub_cache', PRIVATE_CACHE_ROOT)]:
+        if key in st.session_state:
+            try:
+                remove_session_directory(st.session_state[key], root)
+            except OSError:
+                pass  # Files still held by a disconnected worker disappear on restart.
+    for key in list(st.session_state):
+        if key not in {'picker_directory', 'picker_address'}:
+            del st.session_state[key]
+
+
+def activate_model(location, tokenizer_location=None, display_name=None):
+    model = Model(location)
+    try:
+        p = model.path
+        auto = (p if p.is_dir() else p.parent) / 'tokenizer.json'
+        selected_tokenizer = tokenizer_location or (auto if auto.exists() else None)
+        tok = load_tokenizer(selected_tokenizer) if selected_tokenizer else None
+    except Exception:
+        model.close()
+        raise
+    st.session_state.model = model
+    st.session_state.tok = tok
+    st.session_state.loaded = str(model.path)
+    st.session_state.model_display_name = display_name or model.path.name
+    st.session_state.nearest_results = {}
+    st.session_state.clear_hub_token = True
+
+
+title_col, action_col = st.columns([6, 1], vertical_alignment='center')
+if 'model' in st.session_state:
+    title_col.subheader('Token Visualiser')
+else:
+    title_col.title('Token Visualiser')
+if 'model' in st.session_state:
+    action_col.button('Clear model', key='clear_model', on_click=clear_model, use_container_width=True)
+if 'model' not in st.session_state:
+    st.caption('Explore the numbers inside a token. Models run in this Space · No GPU required' if hosted else 'Explore the numbers inside a token. Local model files · No GPU required')
+
+if 'model' not in st.session_state:
+    _, loader, _ = st.columns([1, 3, 1])
+    with loader, st.container(border=True):
+        st.header('Load a model')
+        st.write('Open a local model or download weights from Hugging Face to explore token embeddings.')
+        if hosted:
+            hub_tab, local_tab = st.tabs(['Hugging Face', 'Upload files'])
+        else:
+            local_tab, hub_tab = st.tabs(['Local files', 'Hugging Face'])
+        with local_tab:
+            if hosted:
+                render_upload_loader(activate_model)
+        if not hosted:
+            with local_tab:
+                if st.button('Open model…', key='open_model', use_container_width=True):
+                    st.session_state.picker_open = True
+                if st.session_state.get('picker_open', False):
+                    show_model_picker()
+                location = st.text_input('Model file or folder', placeholder=r'C:\Models\gemma.gguf', key='model_location')
+                with st.expander('Tokenizer (optional)'):
+                    tokenizer_path = st.text_input('Tokenizer JSON or folder', key='tokenizer_location',
+                        help='Use the tokenizer from the exact same model. Enables text tokenization without added special tokens. Detected automatically beside the model when available.')
+                load = st.button('Load model', key='load_model', type='primary', disabled=not location.strip(), use_container_width=True)
+                st.caption('GGUF · F32/F16/BF16 safetensors · Sharded model folders')
+                st.caption('Files stay on this computer.')
+                if load:
+                    try:
+                        with st.spinner('Reading model metadata…'):
+                            activate_model(location.strip().strip('"'), tokenizer_path)
+                        st.rerun()
+                    except Exception as e:
+                        st.error(f'Could not load model: {e}')
+        with hub_tab:
+            render_hub_loader(activate_model)
+    st.stop()
+
+model = st.session_state.model
+st.markdown('<style>' + Path(__file__).with_name('viewer.css').read_text(encoding='utf-8') + '</style>', unsafe_allow_html=True)
+chart_pane = st.container(key='chart_pane', height='stretch', border=True)
+chart_slot = chart_pane.empty()
+with st.container(key='controls_pane'):
+    with st.container(key='control_columns'):
+        token_panel, neighbour_panel, calculation_panel, details_panel = st.columns([1.15, 1.6, 1.15, 1.05], gap='medium')
+
+with details_panel:
+    st.subheader('Chart & model')
+    st.caption(st.session_state.get('model_display_name', Path(st.session_state.loaded).name))
+    names = sorted(model.tensors, key=lambda n: (not embedding_name(n), n))
+    name = st.selectbox('Embedding tensor', names)
+    if not embedding_name(name):
+        st.warning('This tensor may not contain input token embeddings.')
+    chart_view = st.selectbox('Chart view', ['Dimensions', 'Relationship 2D', 'Relationship 3D'], key='chart_view')
+    style = st.radio('Chart style', ['Line', 'Scatter'], horizontal=True)
+shape = model.tensors[name][2]
+vocab = model.vocab
+tok = st.session_state.tok
+get_row = lambda i: model.row(name, i)
+read_rows = lambda start, stop: model.rows(name, start, stop)
+with details_panel:
+    a, b = st.columns(2)
+    a.metric('Vocabulary rows', f'{shape[0]:,}')
+    b.metric('Values per token', f'{shape[1]:,}')
+    st.caption('Stored input embeddings, before layer processing. Drag the chart to zoom; double-click to reset. Click a legend entry to hide or show a token.')
+
+with token_panel:
+    ids = []
+    labels = {}
+    comparison_scope = f'{st.session_state.loaded}:{name}'
+    calculation_key = f'calculations:{comparison_scope}'
+    st.subheader('Compare tokens')
+    st.caption('Search, select and add tokens. Each curve has its own colour.')
+    if tok is not None:
+        text = st.text_input('Type text', 'hello world')
+        try:
+            encoded = tok.encode(text, add_special_tokens=False)
+            st.caption('Tokenized pieces: ' + ' · '.join(f'{piece!r} (ID {i})' for piece, i in zip(encoded.tokens, encoded.ids)))
+            choices = list(range(len(encoded.ids)))
+            positions = st.multiselect('Select tokens to plot', choices, default=choices[:2],
+                                       format_func=lambda p: f'{p + 1}: {encoded.tokens[p]!r} · ID {encoded.ids[p]}')
+            for p in positions:
+                i = encoded.ids[p]
+                if i not in ids:
+                    ids.append(i)
+                labels[i] = f'{encoded.tokens[p]} · {i}'
+        except Exception as e:
+            st.error(f'Tokenization failed: {e}')
+    elif vocab is not None:
+        query = st.text_input('Find a vocabulary token', 'hello', help='Searches stored token spellings, not text tokenization. Spaces may appear as ▁ or Ġ.')
+        matches = [i for i, word in enumerate(vocab) if query in word][:200] if query else list(range(min(200, len(vocab))))
+        matches.sort(key=lambda i: (vocab[i] != query, len(vocab[i]), i))
+        # The chart selection lives independently of changing search-result widgets.
+        selection_key = f'chart_tokens:{comparison_scope}'
+        if selection_key not in st.session_state:
+            st.session_state[selection_key] = matches[:1]
+        selected = [i for i in st.session_state[selection_key] if 0 <= i < min(shape[0], len(vocab))]
+        result, add = st.columns([2.3, 1.7], vertical_alignment='bottom')
+        candidate = result.selectbox('Matching token', matches, key=f'match:{comparison_scope}:{query}',
+                                     format_func=lambda i: f'{vocab[i]!r} · ID {i}')
+        if add.button('Add to chart', type='primary', disabled=candidate is None or candidate in selected or len(selected) + len(st.session_state.get(calculation_key, [])) >= 20):
+            selected.append(candidate)
+            st.session_state[selection_key] = selected
+        ids = st.multiselect('Tokens on chart', selected, default=selected,
+                            key=f'chart_selection:{comparison_scope}:{tuple(selected)}',
+                            format_func=lambda i: f'{vocab[i]!r} · ID {i}')
+        if ids != selected:
+            st.session_state[selection_key] = ids
+            st.rerun()
+        st.caption('Search for another token and click Add to chart to compare it. Remove a token with the × beside its name.')
+        labels = {i: f'{vocab[i]} · {i}' for i in ids}
+        if not matches:
+            st.info('No matching vocabulary entries. Try a substring or enter an ID below.')
+    else:
+        st.info('Add the matching tokenizer.json to type text. You can inspect any token by ID below.')
+
+    ids = list(ids)
+    manual = st.text_input('Additional token IDs (comma-separated)', placeholder='0, 42, 100', key='manual_ids')
+    try:
+        for part in manual.split(','):
+            if part.strip():
+                i = int(part.strip())
+                if not 0 <= i < shape[0]:
+                    raise ValueError(f'ID {i} is outside this tensor’s row range.')
+                if i not in ids:
+                    ids.append(i)
+                labels.setdefault(i, f'ID {i}')
+    except ValueError as e:
+        st.error(str(e))
+        st.stop()
+    def token_label(i):
+        if vocab is not None and i < len(vocab):
+            return f'{vocab[i]!r} · ID {i}'
+        if tok is not None:
+            piece = tok.id_to_token(i)
+            if piece is not None:
+                return f'{piece!r} · ID {i}'
+        return f'ID {i}'
+
+
+    extra_key = f'added_neighbours:{comparison_scope}'
+    extra_ids = [i for i in st.session_state.get(extra_key, []) if i not in ids and 0 <= i < shape[0]]
+    if extra_ids:
+        kept = st.multiselect('Added neighbours on chart', extra_ids, default=extra_ids,
+                              format_func=token_label, key=f'added_neighbour_selection:{comparison_scope}:{tuple(extra_ids)}')
+        if kept != extra_ids:
+            st.session_state[extra_key] = kept
+            st.rerun()
+        for i in kept:
+            ids.append(i)
+            labels[i] = token_label(i)
+    if not ids and not st.session_state.get(calculation_key):
+        chart_slot.info('Select a token in the controls below to see its embedding values.')
+        st.stop()
+    if len(ids) > 20:
+        st.warning('Select at most 20 tokens for a readable comparison.')
+        st.stop()
+
+
+    def add_neighbour(i):
+        added = list(st.session_state.get(extra_key, []))
+        if i not in added:
+            added.append(i)
+        st.session_state[extra_key] = added
+
+
+try:
+    vectors = {i: get_row(i) for i in ids}
+except Exception as e:
+    chart_slot.error(f'Could not read embedding: {e}')
+    st.stop()
+
+calculations = st.session_state.setdefault(calculation_key, [])
+for calculation in calculations:
+    vectors[calculation['id']] = calculation['values']
+    labels[calculation['id']] = calculation['label']
+
+
+def remove_calculation(calculation_id):
+    st.session_state[calculation_key] = [c for c in st.session_state[calculation_key] if c['id'] != calculation_id]
+
+
+def operand_label(key):
+    return labels[key] if isinstance(key, str) else token_label(key)
+
+
+def expression_label(key):
+    if isinstance(key, str):
+        return '(' + labels[key].removeprefix('calc: ') + ')'
+    if vocab is not None and key < len(vocab):
+        return repr(vocab[key])
+    if tok is not None and tok.id_to_token(key) is not None:
+        return repr(tok.id_to_token(key))
+    return f'ID {key}'
+
+
+def save_calculation(values, label):
+    label = label if label.startswith('calc: ') else 'calc: ' + label
+    serial = st.session_state.get('calculation_serial', 0) + 1
+    while label in labels.values():
+        label += f' [{serial}]'
+    st.session_state.calculation_serial = serial
+    calculations.append({'id': f'calc:{serial}', 'label': label, 'values': np.asarray(values).copy()})
+    st.rerun()
+
+
+with calculation_panel:
+    st.subheader('Vector tools')
+    arithmetic_tab, relationship_tab = st.tabs(['Calculations', 'Relationships'])
+with arithmetic_tab:
+    st.caption('Calculate A − B or A + B, dimension by dimension. Results become derived curves labelled calc:.')
+    operands = list(vectors)
+    left_operand = st.selectbox('Vector A', operands, format_func=operand_label, key=f'calc_a:{comparison_scope}')
+    operation = st.selectbox('Operation', ['Subtract (A − B)', 'Add (A + B)'], key=f'calc_operation:{comparison_scope}')
+    right_operand = st.selectbox('Vector B', operands, index=min(1, len(operands) - 1),
+                                 format_func=operand_label, key=f'calc_b:{comparison_scope}')
+    result_name = st.text_input('Result name (optional)', placeholder='e.g. hello minus world', key=f'calc_name:{comparison_scope}')
+    if st.button('Calculate & add to chart', key='add_calculation', type='primary',
+                 disabled=len(vectors) >= 20, use_container_width=True):
+        subtract = operation.startswith('Subtract')
+        values = vectors[left_operand] - vectors[right_operand] if subtract else vectors[left_operand] + vectors[right_operand]
+        symbol = '−' if subtract else '+'
+        expression = f'{expression_label(left_operand)} {symbol} {expression_label(right_operand)}'
+        label = 'calc: ' + (result_name.strip() or expression)
+        save_calculation(values, label)
+    if calculations:
+        st.caption('Calculated curves stay available when you change the selected tokens. They are included in CSV exports.')
+        for calculation in calculations:
+            label_col, remove_col = st.columns([3, 1.2], vertical_alignment='center')
+            label_col.text(calculation['label'])
+            remove_col.button('Remove', key=f'remove_calculation:{calculation["id"]}',
+                              on_click=remove_calculation, args=(calculation['id'],))
+    if len(vectors) >= 20:
+        st.caption('The chart supports up to 20 token and calculated curves in total.')
+
+with relationship_tab:
+    analysis = render_relationship_controls(vectors, operand_label, comparison_scope, save_calculation)
+
+if len(vectors) > 20:
+    token_panel.warning('Remove a token or calculated curve to keep the chart within its 20-curve limit.')
+    st.stop()
+chart_ids = list(vectors)
+colour_key = f'token_colours:{comparison_scope}'
+colours = st.session_state.setdefault(colour_key, {})
+# Retain colours across searches, keeping every currently selected token distinct.
+used = set()
+for i in chart_ids:
+    if i not in colours or colours[i] in used:
+        colours[i] = next(c for c in qualitative.Light24 if c not in used and
+                         c not in {colours.get(other) for other in chart_ids if other != i})
+    used.add(colours[i])
+details_panel.caption(f'Tokens: {len(ids)} · Calculations: {len(calculations)}')
+if chart_view == 'Dimensions':
+    fig = go.Figure()
+    for i, values in vectors.items():
+        fig.add_trace(go.Scattergl(x=np.arange(len(values)), y=values,
+                                  mode='lines' if style == 'Line' else 'markers', name=labels[i],
+                                  line=dict(color=colours[i], width=2 if isinstance(i, str) else 1.5,
+                                            dash='dash' if isinstance(i, str) else 'solid'),
+                                  marker=dict(color=colours[i], size=4),
+                                  hovertemplate='Dimension %{x}<br>Value %{y:.7f}<extra>%{fullData.name}</extra>'))
+    fig.update_layout(autosize=True, xaxis_title='Embedding dimension (zero-based)', yaxis_title='Stored embedding value',
+                      template='plotly_dark', hovermode='x', margin=dict(l=20, r=20, t=20, b=30),
+                      legend=dict(orientation='h', y=1.02, yanchor='bottom', title_text='Tokens'), showlegend=True,
+                      uirevision=f'{comparison_scope}:{style}')
+    fig.add_hline(y=0, line_width=1, line_color='gray')
+    chart_slot.plotly_chart(fig, width='stretch', height='stretch', config={'scrollZoom': True}, key='embedding_chart')
+
+elif analysis is None:
+    chart_slot.info('Open Relationships in Vector tools, choose your inputs, then click Show relationship.')
+else:
+    fig = relationship_figure(analysis, 2 if chart_view == 'Relationship 2D' else 3, vectors, labels)
+    chart_slot.plotly_chart(fig, width='stretch', height='stretch', key='embedding_chart', config={'scrollZoom': True})
+
+with neighbour_panel:
+    st.subheader('Closest tokens')
+    search_mode = st.selectbox('Search mode', ['Nearest tokens', 'Analogy answers'], key=f'neighbour_mode:{comparison_scope}')
+    results = st.session_state.setdefault('nearest_results', {})
+    eligible = True
+    report_id = None
+    if search_mode == 'Nearest tokens':
+        source_id = st.selectbox('Find neighbours of', chart_ids, format_func=operand_label,
+                                 key=f'neighbour_source:{comparison_scope}')
+        query_vector = vectors[source_id]
+        excluded = [source_id] if isinstance(source_id, int) else []
+        result_key = ('ranked', comparison_scope, source_id)
+        row_key = str(source_id)
+        search_label = 'Find closest tokens'
+        st.caption('Cosine similarity uses the full vector, including calculated queries.')
+    elif analysis is None:
+        st.info('Show a relationship in Vector tools before searching for analogy answers.')
+        eligible = False
+    else:
+        query_vector = analysis['metrics']['prediction']
+        excluded = analysis['exclude_ids']
+        expected = analysis['expected']
+        report_id = expected['key'] if expected and isinstance(expected['key'], int) else None
+        result_key = ('analogy', comparison_scope, analysis['id'])
+        row_key = f'analogy:{analysis["id"]}'
+        search_label = 'Refresh answers'
+        st.caption('Search near C + mean(B − A), excluding the input tokens. Results refer to the last shown relationship.')
+    if eligible:
+        count_col, button_col = st.columns([1, 2], vertical_alignment='bottom')
+        neighbour_count = count_col.number_input('Results', min_value=1, max_value=50, value=10, step=1)
+        search = button_col.button(search_label, type='primary', key='find_neighbours')
+        requested = st.session_state.pop(f'run_analogy:{comparison_scope}', False) if search_mode == 'Analogy answers' else False
+        previous_search = any(isinstance(key, tuple) and len(key) == 3 and key[:2] == (comparison_scope, source_id)
+                              for key in results) if search_mode == 'Nearest tokens' else False
+        if search or requested or (result_key not in results and previous_search):
+            bar = st.progress(0, text='Scanning embedding vectors…')
+            try:
+                found = search_embeddings(read_rows, shape[0], query_vector, count=50, exclude_ids=excluded, report_id=report_id,
+                    progress=lambda done, total: bar.progress(done / total, text=f'Scanning embeddings: {done:,} / {total:,}'))
+                results[result_key] = found if search_mode == 'Analogy answers' else found['neighbours']
+            except Exception as e:
+                st.error(f'Could not find neighbours: {e}')
+            finally:
+                bar.empty()
+        if result_key in results:
+            found = results[result_key]
+            if search_mode == 'Analogy answers':
+                neighbours = found['neighbours'][:int(neighbour_count)]
+                if report_id is not None:
+                    if found['expected_rank'] is not None:
+                        st.caption(f'Expected {token_label(report_id)}: rank {found["expected_rank"]:,} of {found["valid_count"]:,}, cosine {found["expected_score"]:.6f}.')
+                    else:
+                        st.caption('Expected D is excluded as an input token, or has a zero/nonfinite embedding.')
+            else:
+                neighbours = found[:int(neighbour_count)]
+            if neighbours:
+                st.caption(f'Showing {len(neighbours)} neighbours. Change Results to show more or fewer without searching again.')
+                with st.container(border=False):
+                    header = st.columns([.35, 1.7, 1, 1.8])
+                    for column, title in zip(header, ['#', 'Token', 'Cosine', '']):
+                        if title:
+                            column.markdown(f'**{title}**')
+                    for rank, (i, score) in enumerate(neighbours, 1):
+                        rank_col, token_col, score_col, action_col = st.columns([.35, 1.7, 1, 1.8], vertical_alignment='center')
+                        rank_col.text(str(rank))
+                        token_col.text(token_label(i))
+                        score_col.write(f'{score:.6f}')
+                        action_col.button('On chart' if i in ids else 'Add to chart',
+                            key=f'add_neighbour:{comparison_scope}:{row_key}:{i}',
+                            disabled=i in ids or len(vectors) >= 20, on_click=add_neighbour, args=(i,), use_container_width=True)
+                if len(vectors) >= 20:
+                    st.caption('The chart has reached its 20-curve limit. Remove a curve before adding another.')
+                st.caption('Cosine similarity ranges from −1 to 1. Nearby tokens can reflect spelling, boundaries or usage as well as meaning.')
+            else:
+                st.info('No other eligible tokens with finite, nonzero embeddings were found.')
+with details_panel:
+    st.dataframe([{'Token': labels[i], 'Min': float(v.min()), 'Max': float(v.max()),
+                   'Mean': float(v.mean()), 'Std dev': float(v.std()), 'L2 norm': float(np.linalg.norm(v))}
+                  for i, v in vectors.items()], use_container_width=True, hide_index=True)
+    output = io.StringIO()
+    writer = csv.writer(output)
+    writer.writerow(['dimension'] + [labels[i] for i in chart_ids])
+    writer.writerows([d] + [float(vectors[i][d]) for i in chart_ids] for d in range(shape[1]))
+    st.download_button('Download values as CSV', output.getvalue(), 'token_embeddings.csv', 'text/csv')
+    with st.expander('Inspect exact values'):
+        st.dataframe({'dimension': np.arange(shape[1]), **{labels[i]: vectors[i] for i in chart_ids}}, use_container_width=True)