import functools

import os

import tkinter as tk

import threading

from collections import defaultdict, ChainMap

from functools import reduce

from pprint import pprint

from tkinter import filedialog, ttk

from tkinter import messagebox

from tkinter.font import Font

import re

import PIL

import pyperclip

import bisect

import traceback

from view.colors import history_color, not_visited_color, visited_color, ooc_color, text_color, uncanonical_color

from view.display import Display

from view.dialogs import GenerationSettingsDialog, InfoDialog, VisualizationSettingsDialog, \

NodeChapterDialog, MultimediaDialog, NodeInfoDialog, SearchDialog, GotoNode, \

PreferencesDialog, AIMemory, CreateMemory, NodeMemory, ChatSettingsDialog, CreateSummary, Summaries

from model import TreeModel

from util.util import clip_num, metadata, diff

from util.util_tree import depth, height, flatten_tree, stochastic_transition, node_ancestry, subtree_list, node_index, \

nearest_common_ancestor

from util.gpt_util import logprobs_to_probs

def gated_call(f, condition):

def _gated_call(*args, _f=f, _cond=condition, **kwargs):

if _cond():

# print('cond')

_f(*args, **kwargs)

# print("no cond")

return _gated_call

def no_junk_args(f):

return lambda event=None, *args, _f=f, **kwargs: _f(*args, **kwargs)

class Controller:

def __init__(self, root):

self.callbacks = self.build_callbacks()

self.root = root

self.state = TreeModel(self.root)

self.display = Display(self.root, self.callbacks, self.state, self)

self.register_model_callbacks()

self.setup_key_bindings()

self.build_menus()

self.ancestor_end_indices = None

# move to preferences dict in state

self.canonical_only = False

self.display_canonical = True

#################################

# Hooks

#################################

def build_callbacks(self):

attrs = [getattr(self, f) for f in dir(self)]

callbacks = [f for f in attrs if callable(f) and hasattr(f, "meta") and "name" in f.meta]

return {

f.meta["name"]: {**f.meta, "callback": no_junk_args(f)}

for f in callbacks

}

def register_model_callbacks(self):

# When the tree is updated, refresh the navtree, nav selection, and textbox

self.state.register_callback(self.state.tree_updated, self.update_nav_tree)

self.state.register_callback(self.state.tree_updated, self.update_nav_tree_selected)

self.state.register_callback(self.state.tree_updated, self.update_chapter_nav_tree)

self.state.register_callback(self.state.tree_updated, self.update_chapter_nav_tree_selected)

self.state.register_callback(self.state.tree_updated, self.save_edits)

self.state.register_callback(self.state.tree_updated, self.refresh_textbox)

self.state.register_callback(self.state.tree_updated, self.refresh_visualization)

self.state.register_callback(self.state.tree_updated, self.refresh_display)

# TODO autosaving takes too long for a big tree

self.state.register_callback(self.state.tree_updated, lambda **kwargs: self.save_tree(popup=False, autosave=True))

# Before the selection is updated, save edits

self.state.register_callback(self.state.pre_selection_updated, self.save_edits)

# When the selection is updated, refresh the nav selection and textbox

self.state.register_callback(self.state.selection_updated, self.update_nav_tree_selected)

self.state.register_callback(self.state.selection_updated, self.update_chapter_nav_tree_selected)

self.state.register_callback(self.state.selection_updated, self.refresh_textbox)

self.state.register_callback(self.state.selection_updated, self.refresh_vis_selection)

self.state.register_callback(self.state.selection_updated, self.refresh_notes)

self.state.register_callback(self.state.selection_updated, self.refresh_counterfactual_meta)

self.state.register_callback(self.state.selection_updated, self.refresh_display)

self.state.register_callback(self.state.io_update, self.update_dropdown)

def setup_key_bindings(self):

attrs = [getattr(self, f) for f in dir(self)]

funcs_with_keys = [f for f in attrs if callable(f) and hasattr(f, "meta") and "keys" in f.meta]

def in_edit():

return self.display.mode in ["Edit", "Child Edit"] \

or (self.display.mode == "Visualize" and self.display.vis.textbox) \

or self.has_focus(self.display.input_box)

valid_keys_outside_edit = ["Control", "Alt", "Escape", "Delete"]

for f in funcs_with_keys:

for tk_key in f.meta["keys"]:

inside_edit = any(v in tk_key for v in valid_keys_outside_edit)

self.root.bind(

tk_key, no_junk_args(f if inside_edit else gated_call(f, lambda: not in_edit()))

)

# Numbers to select children

for i in range(1, 11):

i = i % 10

f = lambda _i=i: self.state.select_child(_i-1)

self.root.bind(

f"<Key-{i}>", no_junk_args(gated_call(f, lambda: not in_edit()))

)

def build_menus(self):

# Tuple of 4 things: Name, Hotkey display text, tkinter key to bind to, function to call (without arguments)

menu_list = {

"View": [

('Toggle visualize mode', 'J', None, no_junk_args(self.toggle_visualization_mode)),

('Visualization settings', 'Ctrl+U', None, no_junk_args(self.visualization_settings_dialog)),

('Collapse node', 'Ctrl-?', None, no_junk_args(self.collapse_node)),

('Collapse subtree', 'Ctrl-minus', None, no_junk_args(self.collapse_subtree)),

('Collapse all except subtree', 'Ctrl-:', None, no_junk_args(self.collapse_all_except_subtree)),

('Expand children', 'Ctrl-\"', None, no_junk_args(self.expand_children)),

('Expand subtree', 'Ctrl-+', None, no_junk_args(self.expand_subtree)),

('Center view', 'L, Ctrl-L', None, no_junk_args(self.center_view)),

('Reset zoom', 'Ctrl-0', None, no_junk_args(self.reset_zoom)),

],

"Edit": [

('Edit mode', 'Ctrl+E', None, no_junk_args(self.toggle_edit_mode)),

("Create parent", 'Alt-Left', None, no_junk_args(self.create_parent)),

("Change parent", 'Shift-P', None, no_junk_args(self.change_parent)),

("New Child", 'H, Ctrl+H, Alt+Right', None, no_junk_args(self.create_child)),

("New Sibling", 'Alt+Down', None, no_junk_args(self.create_sibling)),

("Merge with parent", 'Shift+Left', None, no_junk_args(self.merge_with_parent)),

("Merge with children", 'Shift+Right', None, no_junk_args(self.merge_with_children)),

("Move up", 'Shift+Up', None, no_junk_args(self.move_up)),

("Move up", 'Shift+Down', None, no_junk_args(self.move_down)),

('Prepend newline', 'N, Ctrl+N', None, no_junk_args(self.prepend_newline)),

('Prepend space', 'Ctrl+Space', None, no_junk_args(self.prepend_space)),

('Copy', 'Ctrl+C', None, no_junk_args(self.copy_text)),

('Delete', 'Backspace', None, no_junk_args(self.delete_node)),

('Delete and reassign children', '', None, no_junk_args(self.delete_node_reassign_children)),

],

"Navigate": [

('Return to root', 'R', None, no_junk_args(self.return_to_root)),

('Save checkpoint', 'Ctrl+T', None, no_junk_args(self.save_checkpoint)),

('Go to checkpoint', 'T', None, no_junk_args(self.goto_checkpoint)),

("Bookmark", "B", None, no_junk_args(self.bookmark)),

("Next Bookmark", "D", None, no_junk_args(self.next_bookmark)),

("Prev Bookmark", "A", None, no_junk_args(self.prev_bookmark)),

("Stochastic walk", "W", None, no_junk_args(self.walk)),

("Edit chapter", "Ctrl+Y", None, no_junk_args(self.chapter_dialog)),

("Search", "Ctrl+F", None, no_junk_args(self.search)),

("Goto node by id", "Ctrl+Shift+G", None, no_junk_args(self.goto_node_dialog)),

],

"Generation": [

('Generation settings', 'Ctrl+shift+p', None, no_junk_args(self.generation_settings_dialog)),

('Chat settings', None, None, no_junk_args(self.chat_settings)),

('Generate', 'G, Ctrl+G', None, no_junk_args(self.generate)),

('View summaries', '', None, no_junk_args(self.view_summaries)),

],

"Memory": [

('AI Memory', 'Ctrl+Shift+M', None, no_junk_args(self.ai_memory)),

('Create memory', 'Ctrl+M', None, no_junk_args(self.add_memory)),

('Node memory', 'Alt+M', None, no_junk_args(self.node_memory)),

],

"Flags": [

("Mark node visited", None, None, lambda: self.set_visited(True)),

("Mark node unvisited", None, None, lambda: self.set_visited(False)),

("Mark subtree visited", None, None, lambda: self.set_subtree_visited(True)),

("Mark subtree unvisited", None, None, lambda: self.set_subtree_visited(False)),

("Mark all visited", None, None, lambda: self.set_all_visited(True)),

("Mark all unvisited", None, None, lambda: self.set_all_visited(False)),

("Mark node as prompt", None, None, lambda: self.set_source('prompt')),

("Mark node as AI completion", None, None, lambda: self.set_source('AI')),

("Mark subtree as prompt", None, None, lambda: self.set_subtree_source('prompt')),

("Mark subtree as AI completion", None, None, lambda: self.set_subtree_source('AI')),

("Mark all as prompt", None, None, lambda: self.set_all_source('prompt')),

("Mark all as AI completion", None, None, lambda: self.set_all_source('AI')),

("Archive", None, None, no_junk_args(self.archive)),

("Unarchive", None, None, no_junk_args(self.unarchive)),

],

"Settings": [

('Chat settings', None, None, no_junk_args(self.chat_settings)),

('Generation settings', 'Ctrl+shift+p', None, no_junk_args(self.generation_settings_dialog)),

('Preferences', 'Ctrl+P', None, no_junk_args(self.preferences))

],

"Info": [

("Tree statistics", "I", None, no_junk_args(self.info_dialog)),

('Multimedia', 'U', None, no_junk_args(self.multimedia_dialog)),

('Node metadata', 'Ctrl+Shift+N', None, no_junk_args(self.node_info_dialogue)),

('Preferences', 'Ctrl+P', None, no_junk_args(self.preferences))

],

}

return menu_list

#################################

# Navigation

#################################

# @metadata(name=, keys=, display_key=)

@metadata(name="Next", keys=["<period>", "<Return>", "<Control-period>"], display_key=">")

def next(self):

if self.state.preferences["canonical_only"]:

self.state.next_canonical()

else:

#self.state.traverse_tree(1)

self.select_node(node=self.state.tree_node_dict[self.state.next_id(1)])

# next.meta = dict(name="Next", keys=["<period>", "<Return>"], display_key=">")

# .meta = dict(name=, keys=, display_key=)

@metadata(name="Prev", keys=["<comma>", "<Control-comma>"], display_key="<",)

def prev(self):

if self.state.preferences["canonical_only"]:

self.state.prev_canonical()

else:

#self.state.traverse_tree(-1)

self.select_node(node=self.state.tree_node_dict[self.state.next_id(-1)])

@metadata(name="Go to parent", keys=["<Left>", "<Control-Left>"], display_key="←")

def parent(self):

#self.state.select_parent()

self.select_node(node=self.state.parent())

@metadata(name="Go to child", keys=["<Right>", "<Control-Right>"], display_key="→")

def child(self):

child_id = self.state.child(0)

if child_id is not None:

self.select_node(node=self.state.tree_node_dict[child_id])

@metadata(name="Go to next sibling", keys=["<Down>", "<Control-Down>"], display_key="↓")

def next_sibling(self):

#self.state.select_sibling(1)

self.select_node(node=self.state.sibling(1))

@metadata(name="Go to previous Sibling", keys=["<Up>", "<Control-Up>"], display_key="↑")

def prev_sibling(self):

#self.state.select_sibling(-1)

self.select_node(node=self.state.sibling(-1))

@metadata(name="Walk", keys=["<Key-w>", "<Control-w>"], display_key="w")

def walk(self, canonical_only=False):

filter_set = self.state.calc_canonical_set() if canonical_only else None

if 'children' in self.state.selected_node and len(self.state.selected_node['children']) > 0:

chosen_child = stochastic_transition(self.state.selected_node, mode='descendents', filter_set=filter_set)

#self.state.select_node(chosen_child['id'])

self.select_node(node=chosen_child)

@metadata(name="Return to root", keys=["<Key-r>", "<Control-r>"], display_key="r")

def return_to_root(self):

#self.state.select_node(self.state.tree_raw_data["root"]["id"])

self.select_node(node=self.state.tree_raw_data["root"])

@metadata(name="Save checkpoint", keys=["<Control-t>"], display_key="ctrl-t")

def save_checkpoint(self, node=None):

if node is None:

node = self.state.selected_node

self.state.checkpoint = node['id']

self.state.tree_updated(edit=[node['id']])

@metadata(name="Go to checkpoint", keys=["<Key-t>"], display_key="t")

def goto_checkpoint(self):

if self.state.checkpoint:

#self.state.select_node(self.state.checkpoint)

self.select_node(node=self.state.tree_node_dict[self.state.checkpoint])

@metadata(name="Nav Select")

def nav_select(self, *, node_id):

if not node_id or node_id == self.state.selected_node_id:

return

if self.change_parent.meta["click_mode"]:

self.change_parent(node=self.state.tree_node_dict[node_id])

# TODO This causes infinite recursion from the vis node. Need to change how updating open status works

# Update the open state of the node based on the nav bar

# node = self.state.tree_node_dict[node_id]

# node["open"] = self.display.nav_tree.item(node["id"], "open")

#self.state.select_node(node_id)

self.select_node(node=self.state.tree_node_dict[node_id])

@metadata(name="Bookmark", keys=["<Key-b>", "<Control-b>"], display_key="b")

def bookmark(self, node=None):

if node is None:

node = self.state.selected_node

node["bookmark"] = not node.get("bookmark", False)

self.state.tree_updated(edit=[node['id']])

@metadata(name="Toggle canonical", keys=["<Control-Shift-KeyPress-C>"], display_key="ctrl+shift+C")

def toggle_canonical(self, node=None):

if node is None:

node = self.state.selected_node

self.state.toggle_canonical(node=node)

self.state.tree_updated(edit=[n['id'] for n in node_ancestry(node, self.state.tree_node_dict)])

@metadata(name="Toggle archived", keys=["<exclam>"], display_key="")

def toggle_archived(self, node=None):

node = node if node else self.state.selected_node

if 'archived' in node and node['archived']:

self.unarchive(node)

else:

self.archive(node)

@metadata(name="Toggle prompt", keys=["<asterisk>"], display_key="")

def toggle_prompt(self, node=None):

self.state.preferences['show_prompt'] = not self.state.preferences['show_prompt']

self.refresh_textbox()

def archive(self, node=None):

node = node if node else self.state.selected_node

node['archived'] = True

if self.state.preferences['hide_archived']:

self.select_node(self.state.parent(node))

self.state.tree_updated(delete=[node['id']])

def unarchive(self, node=None):

node = node if node else self.state.selected_node

node['archived'] = False

if self.state.preferences['hide_archived']:

self.state.tree_updated(add=[node['id']])

@metadata(name="Go to next bookmark", keys=["<Key-d>", "<Control-d>"])

def next_bookmark(self):

book_indices = {idx: d for idx, d in enumerate(self.state.nodes) if d.get("bookmark", False)}

if len(book_indices) < 1:

return

try:

go_to_book = next(i for i, idx in enumerate(book_indices.keys()) if idx > self.state.tree_traversal_idx)

except StopIteration:

go_to_book = 0

#self.state.select_node(list(book_indices.values())[go_to_book]["id"])

self.select_node(list(book_indices.values())[go_to_book])

@metadata(name="Go to prev bookmark", keys=["<Key-a>", "<Control-a>"])

def prev_bookmark(self):

book_indices = {i: d for i, d in enumerate(self.state.nodes) if d.get("bookmark", False)}

if len(book_indices) < 1:

return

earlier_books = list(i for i, idx in enumerate(book_indices.keys()) if idx < self.state.tree_traversal_idx)

go_to_book = earlier_books[-1] if len(earlier_books) > 0 else -1

#self.state.select_node(list(book_indices.values())[go_to_book]["id"])

self.select_node(list(book_indices.values())[go_to_book])

@metadata(name="Center view", keys=["<Key-l>", "<Control-l>"])

def center_view(self):

#self.display.vis.center_view_on_canvas_coords(*self.display.vis.node_coords[self.state.selected_node_id])

self.display.vis.center_view_on_node(self.state.selected_node)

@metadata(name="Select node")

def select_node(self, node):

if self.state.preferences['coloring'] == 'read':

old_node = self.state.selected_node

self.state.select_node(node['id'])

nca_node, index = nearest_common_ancestor(old_node, node, self.state.tree_node_dict)

nca_end_index = self.ancestor_end_indices[index]

self.display.textbox.tag_delete("old")

self.display.textbox.tag_add("old",

"1.0",

f"1.0 + {nca_end_index} chars")

self.display.textbox.tag_config("old", foreground=history_color())

else:

self.state.select_node(node['id'])

#################################

# Node operations

#################################

@metadata(name="New Child", keys=["<h>", "<Control-h>", "<Alt-Right>"], display_key="h",)

def create_child(self, node=None, update_selection=True, toggle_edit=True):

if node is None:

node = self.state.selected_node

child = self.state.create_child(parent=node, update_selection=update_selection)

self.state.node_creation_metadata(child, source='prompt')

if self.display.mode == "Read" and toggle_edit:

self.toggle_edit_mode()

return child

@metadata(name="New Sibling", keys=["<Alt-Down>"], display_key="alt-down")

def create_sibling(self, node=None):

if node is None:

node = self.state.selected_node

sibling = self.state.create_sibling(node=node)

self.state.node_creation_metadata(sibling, source='prompt')

if self.display.mode == "Read":

self.toggle_edit_mode()

@metadata(name="New Parent", keys=["<Alt-Left>"], display_key="alt-left")

def create_parent(self, node=None):

if node is None:

node = self.state.selected_node

parent = self.state.create_parent(node=node)

self.state.node_creation_metadata(parent, source='prompt')

return parent

@metadata(name="Change Parent", keys=["<Shift-P>"], display_key="shift-p", selected_node=None, click_mode=False)

def change_parent(self, node=None, click_mode=False):

if node is None:

node = self.state.selected_node

if self.change_parent.meta["selected_node"] is None:

self.display.change_cursor("fleur")

self.change_parent.meta["selected_node"] = node

self.change_parent.meta["click_mode"] = click_mode

else:

self.display.change_cursor("arrow")

self.state.change_parent(node=self.change_parent.meta["selected_node"], new_parent_id=node["id"])

self.change_parent.meta["selected_node"] = None

self.change_parent.meta["click_mode"] = False

@metadata(name="Merge with Parent", keys=["<Shift-Left>"], display_key="shift-left",)

def merge_parent(self, node=None):

if node is None:

node = self.state.selected_node

self.state.merge_with_parent(node=node)

@metadata(name="Merge with children", keys=["<Shift-Right>"], display_key="shift-right")

def merge_children(self, node=None):

if node is None:

node = self.state.selected_node

self.state.merge_with_children(node=node)

@metadata(name="Move up", keys=["<Shift-Up>"], display_key="shift-up")

def move_up(self, node=None):

if node is None:

node = self.state.selected_node

self.state.shift(node, -1)

@metadata(name="Move down", keys=["<Shift-Down>"], display_key="shift-down")

def move_down(self, node=None):

if node is None:

node = self.state.selected_node

self.state.shift(node, 1)

@metadata(name="Generate", keys=["<g>", "<Control-g>"], display_key="g")

def generate(self, node=None, **kwargs):

if self.display.mode == "Multiverse":

self.propagate_wavefunction()

else:

if node is None:

node = self.state.selected_node

try:

node["open"] = True

self.display.nav_tree.item(node, open=True)

except Exception as e:

print(str(e))

self.state.generate_continuation(node=node, **kwargs)

def propagate_wavefunction(self):

if self.display.mode == "Multiverse":

if self.display.multiverse.active_wavefunction():

active_node = self.display.multiverse.active_info()

start_position = (active_node['x'], active_node['y'])

multiverse, ground_truth = self.state.generate_greedy_multiverse(max_depth=4, prompt=active_node['prefix'],

unnormalized_amplitude=active_node['amplitude'],

ground_truth="",

threshold=0.04,

engine='babbage')

else:

start_position = (0, 0)

multiverse, ground_truth = self.state.generate_greedy_multiverse(max_depth=4, ground_truth="",

threshold=0.04,

engine='babbage')

self.display.multiverse.draw_multiverse(multiverse=multiverse, ground_truth=ground_truth,

start_position=start_position)

# def propagage_wavefunction_realtime(self):

# if self.display.mode == "Multiverse":

# self.display.multiverse.propagate_realtime(prompt=self.state.build_prompt(quiet=True),

# max_depth=3,

# threshold=0.01,

# engine='ada')

@metadata(name="Delete", keys=["<BackSpace>", "<Control-BackSpace>"], display_key="«")

def delete_node(self, node=None, reassign_children=False):

if node is None:

node = self.state.selected_node

if not node or "parent_id" not in node:

return

result = messagebox.askquestion("Delete", "Delete Node?", icon='warning')

if result != 'yes':

return

self.state.delete_node(node=node, reassign_children=reassign_children)

@metadata(name="Delete and reassign children")

def delete_node_reassign_children(self, node=None):

self.delete_node(node=node, reassign_children=True)

@metadata(name="Enter text", keys=["<Control-bar>"], display_key="Ctrl-Return")

def enter_text(self):

self.display.vis.delete_textbox()

@metadata(name="Escape textbox", keys=["<Escape>"], display_key="Escape")

def escape(self):

self.display.vis.delete_textbox(save=False)

if self.display.in_edit_mode:

self.display.set_mode("Read")

self.refresh_textbox()

if self.change_parent.meta["selected_node"]:

self.display.change_cursor("arrow")

self.change_parent.meta["selected_node"] = None

@metadata(name="Edit history", keys=[], display_key="", persistent_id=None)

def edit_history(self, index):

if self.display.mode == "Read":

_, selected_ancestor = self.index_to_ancestor(index)

self.state.select_node(selected_ancestor["id"])

self.toggle_edit_mode()

@metadata(name="Goto history", keys=[], display_key="")

def goto_history(self, index):

if self.display.mode == "Read":

_, selected_ancestor = self.index_to_ancestor(index)

self.nav_select(node_id=selected_ancestor["id"])

@metadata(name="Split node", keys=[], display_key="")

def split_node(self, index, change_selection=True):

if self.display.mode == "Read":

ancestor_index, selected_ancestor = self.index_to_ancestor(index)

negative_offset = self.ancestor_end_indices[ancestor_index] - index

split_index = len(selected_ancestor['text']) - negative_offset

new_parent, _ = self.state.split_node(selected_ancestor, split_index)

if change_selection:

self.nav_select(node_id=new_parent["id"])

# TODO deal with metadata

def index_to_ancestor(self, index):

ancestor_index = bisect.bisect_left(self.ancestor_end_indices, index)

return ancestor_index, node_ancestry(self.state.selected_node, self.state.tree_node_dict)[ancestor_index]

#################################

# Token manipulation

#################################

@metadata(name="Select token", keys=[], display_key="", selected_node=None, token_index=None)

def select_token(self, index):

if self.display.mode == "Read":

self.display.textbox.tag_remove("selected", "1.0", 'end')

ancestor_index, selected_node = self.index_to_ancestor(index)

negative_offset = self.ancestor_end_indices[ancestor_index] - index

offset = len(selected_node['text']) - negative_offset

# TODO new token offsets if changed

if "meta" in selected_node and "generation" in selected_node["meta"] and not selected_node['meta']['modified']:

self.change_token.meta["counterfactual_index"] = 0

self.change_token.meta["prev_token"] = None

#token_offsets = [n - len(selected_node['meta']['generation']['prompt'])

# for n in selected_node['meta']['generation']["logprobs"]["text_offset"]]

token_offsets = selected_node['meta']['generation']["logprobs"]["text_offset"]

token_index = bisect.bisect_left(token_offsets, offset) - 1

start_position = token_offsets[token_index]

token = selected_node['meta']['generation']["logprobs"]["tokens"][token_index]

#print(selected_node['meta']['generation']["logprobs"]["top_logprobs"])

#print(token_index)

counterfactuals = selected_node['meta']['generation']["logprobs"]["top_logprobs"][token_index]

if self.state.preferences['prob']:

counterfactuals = {k: logprobs_to_probs(v) for k, v in sorted(counterfactuals.items(), key=lambda item: item[1], reverse=True)}

self.print_to_debug(counterfactuals)

#print('start position: ', self.ancestor_end_indices[ancestor_index - 1] + start_position)

self.display.textbox.tag_add("selected", f"1.0 + {self.ancestor_end_indices[ancestor_index - 1] + start_position} chars",

f"1.0 + {self.ancestor_end_indices[ancestor_index - 1] + start_position + len(token)} chars")

#self.change_token(selected_node, token_index)

self.select_token.meta["selected_node"] = selected_node

self.select_token.meta["token_index"] = token_index

@metadata(name="Change token", keys=[], display_key="", counterfactual_index=0, prev_token=None, temp_token_offsets=None)

def change_token(self, node=None, token_index=None, traverse=1):

if not self.select_token.meta["selected_node"]:

return

elif not node:

node = self.select_token.meta["selected_node"]

token_index = self.select_token.meta["token_index"]

if not self.change_token.meta['temp_token_offsets']:

#token_offsets = [n - len(node['meta']['generation']['prompt'])

# for n in node['meta']['generation']["logprobs"]["text_offset"]]

token_offsets = node['meta']['generation']["logprobs"]["text_offset"]

self.change_token.meta['temp_token_offsets'] = token_offsets

else:

token_offsets = self.change_token.meta['temp_token_offsets']

start_position = token_offsets[token_index]

token = node['meta']['generation']["logprobs"]["tokens"][token_index]

counterfactuals = node['meta']['generation']["logprobs"]["top_logprobs"][token_index].copy()

original_token = (token, counterfactuals.pop(token, None))

index = node_index(node, self.state.tree_node_dict)

sorted_counterfactuals = list(sorted(counterfactuals.items(), key=lambda item: item[1], reverse=True))

sorted_counterfactuals.insert(0, original_token)

self.change_token.meta["counterfactual_index"] += traverse

if self.change_token.meta["counterfactual_index"] < 0 or self.change_token.meta["counterfactual_index"] > len(sorted_counterfactuals) - 1:

self.change_token.meta["counterfactual_index"] -= traverse

return

new_token = sorted_counterfactuals[self.change_token.meta["counterfactual_index"]][0]

if not self.change_token.meta['prev_token']:

self.change_token.meta['prev_token'] = token

token_start = self.ancestor_end_indices[index - 1] + start_position

self.display.textbox.config(state="normal")

self.display.textbox.delete(f"1.0 + {token_start} chars",

f"1.0 + {token_start + len(self.change_token.meta['prev_token'])} chars")

self.display.textbox.insert(f"1.0 + {token_start} chars", new_token)

self.display.textbox.config(state="disabled")

self.display.textbox.tag_add("modified",

f"1.0 + {token_start} chars",

f"1.0 + {token_start + len(new_token)} chars")

#update temp token offsets

diff = len(new_token) - len(self.change_token.meta['prev_token'])

for index, offset in enumerate(self.change_token.meta['temp_token_offsets'][token_index + 1:]):

self.change_token.meta['temp_token_offsets'][index + token_index + 1] += diff

self.change_token.meta['prev_token'] = new_token

@metadata(name="Next token", keys=["<Alt-period>"], display_key="", counterfactual_index=0, prev_token=None)

def next_token(self, node=None, token_index=None):

self.change_token(node, token_index, traverse=1)

@metadata(name="Prev token", keys=["<Alt-comma>"], display_key="", counterfactual_index=0, prev_token=None)

def prev_token(self, node=None, token_index=None):

self.change_token(node, token_index, traverse=-1)

@metadata(name="Apply counterfactual", keys=["<Alt-Return>"], display_key="", counterfactual_index=0, prev_token=None)

def apply_counterfactual_changes(self):

# TODO apply to non selected nodes

index = node_index(self.state.selected_node, self.state.tree_node_dict)

new_text = self.display.textbox.get(f"1.0 + {self.ancestor_end_indices[index - 1]} chars", "end-1c")

self.state.update_text(node=self.state.selected_node, text=new_text, modified_flag=False)

self.display.textbox.tag_remove("modified", "1.0", 'end-1c')

# TODO don't do this if no meta

self.state.selected_node['meta']['generation']["logprobs"]["text_offset"] = self.change_token.meta['temp_token_offsets']

self.refresh_counterfactual_meta()

@metadata(name="Refresh counterfactual")

def refresh_counterfactual_meta(self):

self.change_token.meta['prev_token'] = None

self.change_token.meta['temp_token_offsets'] = None

#################################

# State

#################################

# Enters edit mode or exits either edit mode

@metadata(name="Edit", keys=["<e>", "<Control-e>"], display_key="e")

def toggle_edit_mode(self, to_edit_mode=None):

if self.display.mode != "Visualize":

self.save_edits()

to_edit_mode = to_edit_mode if to_edit_mode is not None else not self.display.in_edit_mode

if to_edit_mode:

self.display.set_mode("Edit")

else:

if self.edit_history.meta["persistent_id"] is not None:

self.state.select_node(self.edit_history.meta["persistent_id"])

self.edit_history.meta["persistent_id"] = None

self.display.set_mode("Read")

self.refresh_textbox()

else:

if self.display.vis.textbox is None:

self.display.vis.textbox_events[self.state.selected_node['id']]()

else:

self.display.vis.delete_textbox()

@metadata(name="Child Edit", keys=[], display_key="c")

def toggle_child_edit_mode(self, to_edit_mode=None):

self.save_edits()

to_edit_mode = to_edit_mode if to_edit_mode is not None else not self.display.mode == "Multi Edit"

self.display.set_mode("Multi Edit" if to_edit_mode else "Read")

self.refresh_textbox()

@metadata(name="Visualize", keys=["<Key-j>", "<Control-j>"], display_key="j")

def toggle_visualization_mode(self):

if self.state.preferences['autosave']:

self.save_edits()

self.display.set_mode("Visualize" if self.display.mode != "Visualize" else "Read")

self.refresh_visualization()

self.refresh_textbox()

@metadata(name="Multiverse", keys=[])

def toggle_multiverse_mode(self):

if self.state.preferences['autosave']:

self.save_edits()

self.display.set_mode("Multiverse" if self.display.mode != "Multiverse" else "Read")

self.refresh_visualization()

self.refresh_textbox()

self.refresh_display()

#################################

# Edit

#################################

@metadata(name="Merge with parent")

def merge_with_parent(self, node=None):

self.state.merge_with_parent()

@metadata(name="Merge with children")

def merge_with_children(self):

self.state.merge_with_children()

@metadata(name="Copy")

def copy_text(self):

pyperclip.copy(self.display.textbox.get("1.0", "end-1c"))

@metadata(name="Prepend newline", keys=["n", "<Control-n>"], display_key="n")

def prepend_newline(self):

self.save_edits()

if self.state.selected_node:

text = self.state.selected_node["text"]

if text.startswith("\n"):

text = text[1:]

else:

if text.startswith(' '):

text = text[1:]

text = "\n" + text

self.state.update_text(self.state.selected_node, text)

@metadata(name="Prepend space", keys=["<Control-space>"], display_key="ctrl-space")

def prepend_space(self):

self.save_edits()

if self.state.selected_node:

text = self.state.selected_node["text"]

if text.startswith(" "):

text = text[1:]

else:

text = " " + text

self.state.update_text(self.state.selected_node, text)

#################################

# Collapsing

#################################

@metadata(name="Expand subtree", keys=["<Control-plus>"], display_key="Ctrl-plus")

def expand_subtree(self):

pass

@metadata(name="Collapse subtree", keys=["<Control-minus>"], display_key="Ctrl-minus")

def collapse_subtree(self):

pass

@metadata(name="Expand children", keys=["<Control-quotedbl>"], display_key="Ctrl-\"")

def expand_children(self):

pass

@metadata(name="Collapse node", keys=["<Control-question>"], display_key="Ctrl-?")

def collapse_node(self):

pass

@metadata(name="Collapse all except subtree", keys=["<Control-colon>"], display_key="Ctrl-:")

def collapse_all_except_subtree(self):

pass

#################################

# View

#################################

def set_visited(self, status=True):

self.state.selected_node["visited"] = status

self.update_nav_tree()

self.update_nav_tree_selected()

def set_subtree_visited(self, status=True):

for d in flatten_tree(self.state.selected_node):

d["visited"] = status

self.update_nav_tree()

self.update_nav_tree_selected()

def set_all_visited(self, status=True):

for d in self.state.nodes:

d["visited"] = status

self.update_nav_tree()

self.update_nav_tree_selected()

def set_source(self, source='AI', node=None, refresh=True):

if not node:

node = self.state.selected_node

if not 'meta' in node:

node['meta'] = {}

node['meta']['source'] = source

if refresh:

self.refresh_textbox()

self.update_nav_tree()

self.update_nav_tree_selected()

def set_subtree_source(self, source='AI', node=None):

if not node:

node = self.state.selected_node

for d in flatten_tree(node):

self.set_source(source=source, node=d, refresh=False)

self.refresh_textbox()

self.update_nav_tree()

self.update_nav_tree_selected()

def set_all_source(self, source='AI'):

for d in self.state.nodes:

self.set_source(source=source, node=d, refresh=False)

self.refresh_textbox()

self.update_nav_tree()

self.update_nav_tree_selected()

@metadata(name="Toggle prompt", keys=["<Control-Shift-KeyPress-A>"], display_key="")

def toggle_source(self, node=None):

if not node:

node = self.state.selected_node

if 'meta' in node and 'source' in node['meta'] and node['meta']['source'] == 'AI':

self.set_source(source='prompt', node=node)

elif 'meta' in node and 'source' in node['meta'] and node['meta']['source'] == 'prompt':

self.set_source(source='AI', node=node)

# @metadata(name="Darkmode", keys=["<Control-Shift-KeyPress-D>"], display_key="Ctrl-Shift-D")

# def toggle_darkmode(self):

# global darkmode

# darkmode = not darkmode

# print("c",darkmode)

# self.display.frame.pack_forget()

# self.display.frame.destroy()

# self.display = Display(self.root, self.callbacks, self.state)

# self.refresh_textbox()

# self.update_nav_tree()

#################################

# I/O

#################################

@metadata(name="Open", keys=["<o>", "<Control-o>"], display_key="o")

def open_tree(self):

options = {

'initialdir': os.getcwd() + '/data',

'parent': self.root, 'title': "Open a json tree",

'filetypes': [('json files', '.json')]

}

filename = filedialog.askopenfilename(**options)

if not filename:

return

self.state.open_tree(filename)

# TODO repeated code

@metadata(name="Import JSON as subtree", keys=["<Control-Shift-KeyPress-O>"], display_key="ctrl+shift+o")

def import_tree(self):

options = {

'initialdir': os.getcwd() + '/data',

'parent': self.root, 'title': "Import a json tree",

'filetypes': [('json files', '.json')]

}

filename = filedialog.askopenfilename(**options)

if not filename:

return

self.state.import_tree(filename)

@metadata(name="Save", keys=["<s>", "<Control-s>"], display_key="s")

def save_tree(self, popup=True, autosave=False):

if autosave and not self.state.preferences['autosave']:

return

try:

if not autosave and not self.state.preferences['save_counterfactuals']:

self.state.delete_counterfactuals()

self.save_edits()

self.state.save_tree(backup=popup)

if popup:

messagebox.showinfo(title=None, message="Saved!")

except Exception as e:

messagebox.showerror(title="Error", message=f"Failed to Save!\n{str(e)}")

@metadata(name="Save as sibling", keys=["<Alt-e>"], display_key="alt-e")

def save_as_sibling(self):

# TODO fails on root node

if self.display.mode == "Edit":

new_text = self.display.textbox.get("1.0", 'end-1c')

new_active_text = self.display.secondary_textbox.get("1.0", 'end-1c')

self.escape()

sibling = self.state.create_sibling()

self.nav_select(node_id=sibling['id'])

self.state.update_text(sibling, new_text, new_active_text)

@metadata(name="Export to text", keys=["<Control-Shift-KeyPress-X>"], display_key="Ctrl-Shift-X")

def export_text(self):

try:

filename = self.state.tree_filename if self.state.tree_filename \

else os.path.join(os.getcwd() + '/data/text', "export.txt")

filename = filedialog.asksaveasfilename(

initialfile=os.path.splitext(os.path.basename(filename))[0],

initialdir=os.path.dirname(filename),

defaultextension='.txt')

self.state.export_history(self.state.selected_node, filename)

messagebox.showinfo(title=None, message="Exported!")

except Exception as e:

messagebox.showerror(title="Error", message=f"Failed to Export!\n{str(e)}")

@metadata(name="Clear chapters")

def clear_chapters(self):

result = messagebox.askquestion("Clear chapters", "Delete all chapters?", icon='warning')

if result != 'yes':

return

self.state.remove_all_chapters()

@metadata(name="Save As...")

def save_tree_as(self):

filename = self.state.tree_filename if self.state.tree_filename \

else os.path.join(os.getcwd() + '/data', "new_tree.json")

filename = filedialog.asksaveasfilename(

initialfile=os.path.splitext(os.path.basename(filename))[0],

initialdir=os.path.dirname(filename),

defaultextension='.json')

if filename:

self.state.tree_filename = filename

self.save_tree()

return

@metadata(name="Generation Settings", keys=["<Control-p>"], display_key="ctrl-p")

def generation_settings_dialog(self):

dialog = GenerationSettingsDialog(self.display.frame, self.state.generation_settings)

if dialog.result:

#self.save_tree(popup=False)

self.refresh_textbox()

@metadata(name="Visualization Settings", keys=["<Control-u>"], display_key="ctrl-u")

def visualization_settings_dialog(self):

dialog = VisualizationSettingsDialog(self.display.frame, self.state.visualization_settings)

if dialog.result:

#print("Settings saved")

#pprint(self.state.visualization_settings)

self.refresh_visualization()

# self.save_tree(popup=False)

@metadata(name="Show Info", keys=["<i>", "<Control-i>"], display_key="i")

def info_dialog(self):

all_text = "".join([d["text"] for d in self.state.tree_node_dict.values()])

data = {

"Total characters": f'{len(all_text):,}',

"Total words": f'{len(all_text.split()):,}',

"Total pages": f'{len(all_text) / 3000:,.1f}',

"": "",

"Total nodes": f'{len(self.state.tree_node_dict):,}',

"Max depth": height(self.state.tree_raw_data["root"]),

"Max width": max([len(d["children"]) for d in self.state.tree_node_dict.values()])

}

dialog = InfoDialog(self.display.frame, data)

# @metadata(name="Show Chapters", keys=["<Control-t>"], display_key="ctrl-t")

# def chapters_info_dialog(self):

# dialog = ChaptersInfoDialog(self.display.frame, self.state.chapters)

@metadata(name="Chapter settings", keys=["<Control-y>"], display_key="ctrl-y")

def chapter_dialog(self, node=None):

if node is None:

node = self.state.selected_node

dialog = NodeChapterDialog(parent=self.display.frame, node=node, state=self.state)