Add Natural Memory architecture and tooling
This commit is contained in:
WpyQwq
@@ -0,0 +1,307 @@
|
||||
"""Normalize conversation logs into a leak-resistant memory-policy dataset.
|
||||
|
||||
The runtime accepts many local data shapes because real users rarely keep
|
||||
their chat exports in one format. This command converts them to a small,
|
||||
auditable JSONL schema without inventing labels. It understands the current
|
||||
``native_memory`` episode format, the streaming demo format, and a generic
|
||||
format documented in the output manifest.
|
||||
|
||||
The bundled fallback files are bootstrap data for smoke tests. A real user
|
||||
corpus can be supplied with ``--source``/``--eval-source`` and receives the
|
||||
same normalization and group-level split guarantees.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
from collections import Counter, defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable
|
||||
|
||||
|
||||
DEFAULT_SOURCES = (
|
||||
"data/native_memory/train.jsonl",
|
||||
"data/native_memory/eval.jsonl",
|
||||
"data/demo_stream.jsonl",
|
||||
)
|
||||
PROJECT_ROOT = Path(__file__).resolve().parent
|
||||
|
||||
|
||||
def _project_path(value: str | Path) -> Path:
|
||||
path = Path(value)
|
||||
if path.is_absolute() or path.exists():
|
||||
return path
|
||||
return PROJECT_ROOT / path
|
||||
|
||||
|
||||
def _read_jsonl(path: Path) -> Iterable[tuple[int, dict[str, Any]]]:
|
||||
with path.open("r", encoding="utf-8") as handle:
|
||||
for line_number, raw in enumerate(handle, 1):
|
||||
raw = raw.strip()
|
||||
if not raw:
|
||||
continue
|
||||
value = json.loads(raw)
|
||||
if not isinstance(value, dict):
|
||||
raise ValueError(f"{path}:{line_number} must contain a JSON object")
|
||||
yield line_number, value
|
||||
|
||||
|
||||
def _message_text(messages: Any) -> str:
|
||||
if isinstance(messages, str):
|
||||
return messages.strip()
|
||||
if not isinstance(messages, list):
|
||||
return ""
|
||||
parts: list[str] = []
|
||||
for message in messages:
|
||||
if not isinstance(message, dict):
|
||||
continue
|
||||
content = message.get("content", "")
|
||||
if isinstance(content, str) and content.strip():
|
||||
role = str(message.get("role", "user"))
|
||||
parts.append(f"[{role}] {content.strip()}")
|
||||
return "\n".join(parts).strip()
|
||||
|
||||
|
||||
def _user_text(messages: Any) -> str:
|
||||
if isinstance(messages, str):
|
||||
return messages.strip()
|
||||
if not isinstance(messages, list):
|
||||
return ""
|
||||
for message in reversed(messages):
|
||||
if isinstance(message, dict) and message.get("role") == "user":
|
||||
content = message.get("content", "")
|
||||
if isinstance(content, str):
|
||||
return content.strip()
|
||||
return _message_text(messages)
|
||||
|
||||
|
||||
def _explicit_split(path: Path, *, forced: str | None) -> str | None:
|
||||
if forced in {"train", "eval"}:
|
||||
return forced
|
||||
name = path.name.lower()
|
||||
if any(mark in name for mark in ("eval", "valid", "test")):
|
||||
return "eval"
|
||||
if "train" in name:
|
||||
return "train"
|
||||
return None
|
||||
|
||||
|
||||
def _make_example(
|
||||
*,
|
||||
group_id: str,
|
||||
example_id: str,
|
||||
text: str,
|
||||
write_label: float,
|
||||
forget_label: float = 0.0,
|
||||
kind: str = "conversation",
|
||||
source: str,
|
||||
subject: str = "",
|
||||
attribute: str = "",
|
||||
value: Any = None,
|
||||
answer: str = "",
|
||||
answerable: bool | None = None,
|
||||
messages: Any = None,
|
||||
) -> dict[str, Any] | None:
|
||||
text = str(text or "").strip()
|
||||
if not text:
|
||||
return None
|
||||
return {
|
||||
"id": example_id,
|
||||
"group_id": group_id,
|
||||
"text": text,
|
||||
"messages": messages if isinstance(messages, list) else [{"role": "user", "content": text}],
|
||||
"write_label": float(max(0.0, min(1.0, write_label))),
|
||||
"forget_label": float(max(0.0, min(1.0, forget_label))),
|
||||
"kind": kind,
|
||||
"source": source,
|
||||
"subject": str(subject or ""),
|
||||
"attribute": str(attribute or ""),
|
||||
"value": "" if value is None else str(value),
|
||||
"answer": str(answer or ""),
|
||||
"answerable": answerable,
|
||||
}
|
||||
|
||||
|
||||
def normalize_record(record: dict[str, Any], *, source: str, line_number: int) -> list[dict[str, Any]]:
|
||||
"""Convert one source record into labeled write/query decisions."""
|
||||
|
||||
raw_id = str(record.get("id") or record.get("conversation_id") or f"line-{line_number}")
|
||||
group_id = f"{source}:{raw_id}"
|
||||
output: list[dict[str, Any]] = []
|
||||
|
||||
chunks = record.get("memory_chunks")
|
||||
if isinstance(chunks, list):
|
||||
for index, chunk in enumerate(chunks):
|
||||
if not isinstance(chunk, dict):
|
||||
continue
|
||||
messages = chunk.get("messages", [])
|
||||
item = _make_example(
|
||||
group_id=group_id,
|
||||
example_id=f"{raw_id}:memory:{index}",
|
||||
text=_user_text(messages) or str(chunk.get("text", "")),
|
||||
write_label=float(chunk.get("write_label", 1.0)),
|
||||
forget_label=float(chunk.get("forget_label", 0.0)),
|
||||
kind=str(chunk.get("kind", "fact")),
|
||||
source=source,
|
||||
subject=record.get("subject", ""),
|
||||
attribute=record.get("attribute", ""),
|
||||
value=chunk.get("value", record.get("value", "")),
|
||||
messages=messages,
|
||||
)
|
||||
if item is not None:
|
||||
output.append(item)
|
||||
query = record.get("query")
|
||||
query_text = _user_text(query)
|
||||
item = _make_example(
|
||||
group_id=group_id,
|
||||
example_id=f"{raw_id}:query",
|
||||
text=query_text,
|
||||
write_label=0.0,
|
||||
kind="query",
|
||||
source=source,
|
||||
subject=record.get("subject", ""),
|
||||
attribute=record.get("attribute", ""),
|
||||
answer=record.get("answer", ""),
|
||||
answerable=record.get("answerable"),
|
||||
messages=query if isinstance(query, list) else None,
|
||||
)
|
||||
if item is not None:
|
||||
output.append(item)
|
||||
return output
|
||||
|
||||
memory = record.get("memory")
|
||||
if isinstance(memory, list):
|
||||
for index, item_messages in enumerate(memory):
|
||||
item = _make_example(
|
||||
group_id=group_id,
|
||||
example_id=f"{raw_id}:memory:{index}",
|
||||
text=_user_text(item_messages),
|
||||
write_label=1.0,
|
||||
kind="fact",
|
||||
source=source,
|
||||
messages=item_messages if isinstance(item_messages, list) else None,
|
||||
)
|
||||
if item is not None:
|
||||
output.append(item)
|
||||
query = record.get("query")
|
||||
if query is not None:
|
||||
item = _make_example(
|
||||
group_id=group_id,
|
||||
example_id=f"{raw_id}:query",
|
||||
text=_user_text(query),
|
||||
write_label=0.0,
|
||||
kind="query",
|
||||
source=source,
|
||||
answer=record.get("answer", ""),
|
||||
answerable=record.get("answerable"),
|
||||
messages=query if isinstance(query, list) else None,
|
||||
)
|
||||
if item is not None:
|
||||
output.append(item)
|
||||
|
||||
event = record.get("memory_event")
|
||||
if not output and (record.get("text") is not None or record.get("messages") is not None):
|
||||
event = event if isinstance(event, dict) else {}
|
||||
item = _make_example(
|
||||
group_id=group_id,
|
||||
example_id=f"{raw_id}:turn",
|
||||
text=_user_text(record.get("messages")) or str(record.get("text", "")),
|
||||
write_label=float(event.get("write_label", event.get("write", record.get("write_label", 0.0)))),
|
||||
forget_label=float(event.get("forget_label", event.get("forget", record.get("forget_label", 0.0)))),
|
||||
kind=str(event.get("kind", record.get("kind", "conversation"))),
|
||||
source=source,
|
||||
subject=record.get("subject", ""),
|
||||
attribute=record.get("attribute", ""),
|
||||
value=record.get("value", ""),
|
||||
answer=record.get("answer", ""),
|
||||
answerable=record.get("answerable"),
|
||||
messages=record.get("messages"),
|
||||
)
|
||||
if item is not None:
|
||||
output.append(item)
|
||||
return output
|
||||
|
||||
|
||||
def _split_for_group(group_id: str, explicit: str | None, *, eval_ratio: float) -> str:
|
||||
if explicit is not None:
|
||||
return explicit
|
||||
digest = hashlib.sha1(group_id.encode("utf-8")).hexdigest()
|
||||
value = int(digest[:8], 16) / 0xFFFFFFFF
|
||||
return "eval" if value < eval_ratio else "train"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--source", action="append", help="input JSONL; may be repeated")
|
||||
parser.add_argument("--eval-source", action="append", default=[], help="input JSONL forced into eval")
|
||||
parser.add_argument("--output-dir", default="data/production_memory")
|
||||
parser.add_argument("--eval-ratio", type=float, default=0.2)
|
||||
args = parser.parse_args()
|
||||
if not 0.0 < args.eval_ratio < 1.0:
|
||||
raise SystemExit("--eval-ratio must be between 0 and 1")
|
||||
|
||||
source_paths = [_project_path(item) for item in (args.source or DEFAULT_SOURCES)]
|
||||
eval_paths = [_project_path(item) for item in args.eval_source]
|
||||
all_inputs = [(path, None) for path in source_paths] + [(path, "eval") for path in eval_paths]
|
||||
examples: list[tuple[str, dict[str, Any]]] = []
|
||||
source_stats: dict[str, Counter[str]] = defaultdict(Counter)
|
||||
seen: set[tuple[str, str, float, float, str]] = set()
|
||||
for path, forced_split in all_inputs:
|
||||
if not path.exists():
|
||||
raise FileNotFoundError(path)
|
||||
source = str(path)
|
||||
name_split = _explicit_split(path, forced=forced_split)
|
||||
for line_number, record in _read_jsonl(path):
|
||||
normalized = normalize_record(record, source=source, line_number=line_number)
|
||||
for item in normalized:
|
||||
dedupe_key = (
|
||||
item["group_id"],
|
||||
item["text"],
|
||||
item["write_label"],
|
||||
item["forget_label"],
|
||||
item["kind"],
|
||||
)
|
||||
if dedupe_key in seen:
|
||||
source_stats[source]["deduplicated"] += 1
|
||||
continue
|
||||
seen.add(dedupe_key)
|
||||
split = _split_for_group(item["group_id"], name_split, eval_ratio=args.eval_ratio)
|
||||
examples.append((split, item))
|
||||
source_stats[source][split] += 1
|
||||
|
||||
output_dir = _project_path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
split_counts: Counter[str] = Counter()
|
||||
for split in ("train", "eval"):
|
||||
path = output_dir / f"{split}.jsonl"
|
||||
with path.open("w", encoding="utf-8") as handle:
|
||||
for item_split, item in examples:
|
||||
if item_split == split:
|
||||
handle.write(json.dumps(item, ensure_ascii=False) + "\n")
|
||||
split_counts[split] += 1
|
||||
manifest = {
|
||||
"format_version": 1,
|
||||
"schema": {
|
||||
"text": "current turn presented to the write policy",
|
||||
"messages": "optional original chat messages",
|
||||
"write_label": "1 durable memory, 0 ordinary query/casual turn",
|
||||
"forget_label": "1 explicit correction/forget request",
|
||||
"group_id": "conversation/episode identity; never split across train and eval",
|
||||
},
|
||||
"bootstrap_data_warning": "Default files are local bootstrap/synthetic data; pass real exports with --source for production training.",
|
||||
"inputs": [str(path) for path, _ in all_inputs],
|
||||
"counts": dict(split_counts),
|
||||
"source_stats": {key: dict(value) for key, value in source_stats.items()},
|
||||
"dedupe_count": sum(value.get("deduplicated", 0) for value in source_stats.values()),
|
||||
}
|
||||
(output_dir / "manifest.json").write_text(
|
||||
json.dumps(manifest, ensure_ascii=False, indent=2),
|
||||
encoding="utf-8",
|
||||
)
|
||||
print(json.dumps(manifest, ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user