1"""writes indexes to the model directory for the frontend to use or for record keeping"""
  2
  3import importlib.metadata
  4import importlib.resources
  5import inspect
  6import itertools
  7import json
  8from collections.abc import Callable
  9from pathlib import Path
 10
 11import pattern_lens
 12from pattern_lens.attn_figure_funcs import (
 13	_FIGURE_NAMES_KEY,
 14	ATTENTION_MATRIX_FIGURE_FUNCS,
 15)
 16
 17
 18def generate_prompts_jsonl(model_dir: Path) -> None:
 19	"""creates a `prompts.jsonl` file with all the prompts in the model directory
 20
 21	looks in all directories in `{model_dir}/prompts` for a `prompt.json` file
 22	"""
 23	prompts: list[dict] = list()
 24	for prompt_dir in (model_dir / "prompts").iterdir():
 25		prompt_file: Path = prompt_dir / "prompt.json"
 26		if prompt_file.exists():
 27			with open(prompt_file, "r") as f:
 28				prompt_data: dict = json.load(f)
 29				prompts.append(prompt_data)
 30
 31	with open(model_dir / "prompts.jsonl", "w") as f:
 32		for prompt in prompts:
 33			f.write(json.dumps(prompt))
 34			f.write("\n")
 35
 36
 37def generate_models_jsonl(path: Path) -> None:
 38	"""creates a `models.jsonl` file with all the models"""
 39	models: list[dict] = list()
 40	for model_dir in (path).iterdir():
 41		model_cfg_path: Path = model_dir / "model_cfg.json"
 42		if model_cfg_path.exists():
 43			with open(model_cfg_path, "r") as f:
 44				model_cfg: dict = json.load(f)
 45				models.append(model_cfg)
 46
 47	with open(path / "models.jsonl", "w") as f:
 48		for model in models:
 49			f.write(json.dumps(model))
 50			f.write("\n")
 51
 52
 53def get_func_metadata(func: Callable) -> list[dict[str, str | None]]:
 54	"""get metadata for a function
 55
 56	# Parameters:
 57	- `func : Callable` which has a `_FIGURE_NAMES_KEY` (by default `_figure_names`) attribute
 58
 59	# Returns:
 60
 61	`list[dict[str, str | None]]`
 62	each dictionary is for a function, containing:
 63
 64	- `name : str` : the name of the figure
 65	- `func_name : str`
 66		the name of the function. if not a multi-figure function, this is identical to `name`
 67		if it is a multi-figure function, then `name` is `{func_name}.{figure_name}`
 68	- `doc : str` : the docstring of the function
 69	- `figure_save_fmt : str | None` : the format of the figure that the function saves, using the `figure_save_fmt` attribute of the function. `None` if the attribute does not exist
 70	- `source : str | None` : the source file of the function
 71	- `code : str | None` : the source code of the function, split by line. `None` if the source file cannot be read
 72
 73	"""
 74	source_file: str | None = inspect.getsourcefile(func)
 75	output: dict[str, str | None] = dict(
 76		func_name=func.__name__,
 77		doc=func.__doc__,
 78		figure_save_fmt=getattr(func, "figure_save_fmt", None),
 79		source=Path(source_file).as_posix() if source_file else None,
 80	)
 81
 82	try:
 83		output["code"] = inspect.getsource(func)
 84	except OSError:
 85		output["code"] = None
 86
 87	fig_names: list[str] | None = getattr(func, _FIGURE_NAMES_KEY, None)
 88	if fig_names:
 89		return [
 90			{
 91				"name": func_name,
 92				**output,
 93			}
 94			for func_name in fig_names
 95		]
 96	else:
 97		return [
 98			{
 99				"name": func.__name__,
100				**output,
101			},
102		]
103
104
105def generate_functions_jsonl(path: Path) -> None:
106	"unions all functions from `figures.jsonl` and `ATTENTION_MATRIX_FIGURE_FUNCS` into the file"
107	figures_file: Path = path / "figures.jsonl"
108	existing_figures: dict[str, dict] = dict()
109
110	if figures_file.exists():
111		with open(figures_file, "r") as f:
112			for line in f:
113				func_data: dict = json.loads(line)
114				existing_figures[func_data["name"]] = func_data
115
116	# Add any new functions from ALL_FUNCTIONS
117	new_functions_lst: list[dict] = list(
118		itertools.chain.from_iterable(
119			get_func_metadata(func) for func in ATTENTION_MATRIX_FIGURE_FUNCS
120		),
121	)
122	new_functions: dict[str, dict] = {func["name"]: func for func in new_functions_lst}
123
124	all_functions: list[dict] = list(
125		{
126			**existing_figures,
127			**new_functions,
128		}.values(),
129	)
130
131	with open(figures_file, "w") as f:
132		for func_meta in sorted(all_functions, key=lambda x: x["name"]):
133			json.dump(func_meta, f)
134			f.write("\n")
135
136
137def write_html_index(path: Path) -> None:
138	"""writes an index.html file to the path"""
139	html_index: str = (
140		importlib.resources.files(pattern_lens)
141		.joinpath("frontend/index.html")
142		.read_text(encoding="utf-8")
143	)
144	pattern_lens_version: str = importlib.metadata.version("pattern-lens")
145	html_index = html_index.replace("$$PATTERN_LENS_VERSION$$", pattern_lens_version)
146	with open(path / "index.html", "w", encoding="utf-8") as f:
147		f.write(html_index)