-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathbuilder.py
More file actions
138 lines (114 loc) · 5 KB
/
Copy pathbuilder.py
File metadata and controls
138 lines (114 loc) · 5 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
from __future__ import annotations
import importlib
import importlib.util
import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable
from .config import load_json, resolve_path
from .ingest import write_jsonl
from .schema import ChatExample, IntermediateExample, Message
from .sanitizers import apply_sanitizers, build_sanitizer_pipeline
@dataclass(frozen=True)
class BuildResult:
dataset_path: Path
metadata_path: Path
example_count: int
def _load_symbol(spec: str) -> Any:
module_spec, symbol_name = spec.split(":", 1)
if module_spec.endswith(".py"):
module_path = Path(module_spec).resolve()
import_spec = importlib.util.spec_from_file_location(module_path.stem, module_path)
if import_spec is None or import_spec.loader is None:
raise ImportError(f"Could not load module from path: {module_path}")
module = importlib.util.module_from_spec(import_spec)
import_spec.loader.exec_module(module)
else:
module = importlib.import_module(module_spec)
return getattr(module, symbol_name)
def _load_default_system_prompt(config_dir: Path, config: dict[str, Any]) -> str | None:
prompt_path = resolve_path(config_dir, config.get("default_system_prompt_path"))
if prompt_path is None or not prompt_path.exists():
return None
text = prompt_path.read_text(encoding="utf-8").strip()
return text or None
def _normalize_example(
example: ChatExample | IntermediateExample,
default_system_prompt: str | None,
) -> ChatExample:
if isinstance(example, ChatExample):
return ChatExample(conversations=list(example.conversations), metadata=dict(example.metadata))
messages: list[Message] = []
system_prompt = example.system or default_system_prompt
if system_prompt:
messages.append(Message(role="system", content=system_prompt))
messages.append(Message(role="user", content=example.user))
messages.append(Message(role="assistant", content=example.assistant))
return ChatExample(conversations=messages, metadata=dict(example.metadata))
def _read_rows(path: Path) -> list[dict[str, Any]]:
with path.open("r", encoding="utf-8") as handle:
return [json.loads(line) for line in handle if line.strip()]
def _dedupe_chat_examples(
chat_examples: list[ChatExample],
*,
dedupe_user_prompts: bool,
) -> list[ChatExample]:
deduped: list[ChatExample] = []
seen_rows: set[str] = set()
seen_users: set[str] = set()
for example in chat_examples:
row_key = json.dumps(example.to_dataset_row(), sort_keys=True, ensure_ascii=True)
if row_key in seen_rows:
continue
if dedupe_user_prompts:
user_key = example.conversations[1].content.strip()
if user_key in seen_users:
continue
seen_users.add(user_key)
seen_rows.add(row_key)
deduped.append(example)
return deduped
def _assign_example_ids(chat_examples: list[ChatExample]) -> list[ChatExample]:
result: list[ChatExample] = []
for index, example in enumerate(chat_examples, start=1):
metadata = dict(example.metadata)
metadata["example_id"] = f"example-{index:06d}"
result.append(ChatExample(conversations=list(example.conversations), metadata=metadata))
return result
def build_dataset_from_config(config_path: Path) -> BuildResult:
config_path = config_path.resolve()
config_dir = config_path.parent
config = load_json(config_path)
adapter_ctor = _load_symbol(config["adapter"])
adapter = adapter_ctor()
adapter_config = dict(config.get("adapter_config", {}))
sanitizer_pipeline = build_sanitizer_pipeline(config.get("sanitizers", []), loader=_load_symbol)
default_system_prompt = _load_default_system_prompt(config_dir, config)
raw_examples = list(adapter.build_examples(adapter_config, config_dir=config_dir))
sanitized_examples = apply_sanitizers(
raw_examples,
sanitizer_pipeline,
config_dir=config_dir,
)
chat_examples = [
_normalize_example(example, default_system_prompt=default_system_prompt)
for example in sanitized_examples
]
dedupe_config = dict(config.get("dedupe", {}))
if dedupe_config.get("enabled", True):
chat_examples = _dedupe_chat_examples(
chat_examples,
dedupe_user_prompts=dedupe_config.get("dedupe_user_prompts", True),
)
chat_examples = _assign_example_ids(chat_examples)
dataset_path = resolve_path(config_dir, config["output_dataset_path"])
metadata_path = resolve_path(config_dir, config["output_metadata_path"])
assert dataset_path is not None
assert metadata_path is not None
write_jsonl((example.to_dataset_row() for example in chat_examples), dataset_path)
write_jsonl((example.to_metadata_row() for example in chat_examples), metadata_path)
return BuildResult(
dataset_path=dataset_path,
metadata_path=metadata_path,
example_count=len(chat_examples),
)