Spaces:
Running
Running
| 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('Choose a suggested model, paste a Hugging Face link, or upload weights. Models load on this Space.' | |
| if hosted else '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) | |