Upload prep.py with huggingface_hub

prep.py

@@ -78,13 +78,16 @@ FILE_Q = [
     "Read the file {path} and tell me what is on the first line.",
     "What is inside {path}?", "List the files in {dir}.",
 ]
-SYSTEM = ("You are clanker, a helpful assistant that can THINK and USE TOOLS. "
-          "When you need to compute or inspect something, wrap a tool call in "
-          "<tool name=\"...\">arguments</tool>. Available tools: calc(expr), "
-          "python(code), read_file(path), list_dir(path), retrieve(query). After "
-          "a tool result appears in <result>...</result>, continue and give the "
-          "final answer. If <context>...</context> is provided, use it. You may "
-          "use <think>...</think> to reason first.")
+SYSTEM = ("You are clanker, a helpful assistant that can THINK, USE TOOLS, and "
+           "USE MEMORY. You may reason in <think>...</think> at any point, "
+           "interleaved with actions. Wrap tool calls in "
+           "<tool name=\"...\">arguments</tool>. Available tools: calc(expr), "
+           "python(code), read_file(path), list_dir(path), retrieve(query). After "
+           "a tool result appears in <result>...</result>, continue and give the "
+           "final answer. If <context>...</context> is provided, use it. You keep "
+           "facts in a secondary memory: <mem_write>KEY<mem_kv>VALUE</mem_kv> to "
+           "store, <mem_read>KEY</mem_read> to recall (result returns inside "
+           "<mem_kv>...</mem_kv>), and <mem_evict>KEY</mem_evict> to forget.")
 
 
 def gen_synthetic(n=60000):
@@ -135,6 +138,85 @@ def gen_synthetic(n=60000):
     return out
 
 
+# --------------------------------------------------------------------------
+# 2b) Interleaved-thinking + learned-memory examples.
+#     Teaches the model to reason step-by-step *between* tool calls and to
+#     persist/recall facts via <mem_write>/<mem_read>/<mem_evict> against the
+#     side store (see memstore.py). Thinking is INTERLEAVED: think, act,
+#     think, answer -- not one big block at the start.
+# --------------------------------------------------------------------------
+MEM_FACTS = [
+    ("project:clanker", "clanker is a hybrid AR/diffusion LM with learned memory."),
+    ("user:name", "The user's name is Ada."),
+    ("user:likes", "The user likes concise answers and tool use."),
+    ("fact:pi", "pi is approximately 3.14159."),
+    ("fact:capitals", "The capital of France is Paris; of Japan is Tokyo."),
+    ("pref:format", "Prefer <think> reasoning before tool calls."),
+]
+MEM_KEYS = [k for k, _ in MEM_FACTS]
+
+
+def gen_memory(n=30000):
+    out = []
+    for _ in range(n):
+        key, val = random.choice(MEM_FACTS)
+        mode = random.random()
+        if mode < 0.4:
+            # write then read back (persistence demo)
+            conv = (f"<bos><system>{SYSTEM}</system>"
+                    f"<user>Remember that {val}</user>"
+                    f"<assistant><think>I should store this in secondary memory "
+                    f"so I can recall it later.</think>"
+                    f"<mem_write>{key}<mem_kv>{val}</mem_kv>"
+                    f"<think>Stored. Now I can read it back to confirm.</think>"
+                    f"<mem_read>{key}</mem_read><mem_kv>{val}</mem_kv>"
+                    f"Got it -- I'll remember {val}</assistant><eos>")
+        elif mode < 0.75:
+            # read an existing fact, interleaved with reasoning
+            conv = (f"<bos><system>{SYSTEM}</system>"
+                    f"<user>What do you know about {key}?</user>"
+                    f"<assistant><think>Let me pull this from secondary memory.</think>"
+                    f"<mem_read>{key}</mem_read><mem_kv>{val}</mem_kv>"
+                    f"<think>That matches what I stored.</think> "
+                    f"Based on memory: {val}</assistant><eos>")
+        else:
+            # evict
+            conv = (f"<bos><system>{SYSTEM}</system>"
+                    f"<user>Forget {key}.</user>"
+                    f"<assistant><think>I'll remove it from secondary memory.</think>"
+                    f"<mem_evict>{key}</mem_evict>Done, I forgot {key}.</assistant><eos>")
+        out.append(conv)
+    return out
+
+
+# Multi-step reasoning with INTERLEAVED think/act/think/answer.
+REASON_QA = [
+    ("A train travels 60 km/h for 2 hours, then 90 km/h for 1 hour. Total distance?",
+     "60*2 + 90*1 = 120 + 90 = 210 km"),
+    ("If I buy 3 items at $4.50 each and a $2 tax, total cost?",
+     "3*4.50 = 13.50; +2 = 15.50"),
+    ("A rectangle is 8 by 5. Area and perimeter?",
+     "area 8*5=40; perimeter 2*(8+5)=26"),
+    ("Compound 5% on $1000 for 2 years?",
+     "1000*1.05^2 = 1102.50"),
+    ("Mix 2L at 10C with 3L at 40C, final temp?",
+     "(2*10+3*40)/5 = 140/5 = 28C"),
+]
+def gen_interleaved(n=30000):
+    out = []
+    for _ in range(n):
+        q, ans = random.choice(REASON_QA)
+        conv = (f"<bos><system>{SYSTEM}</system>"
+                f"<user>{q}</user>"
+                f"<assistant><think>Break it into parts.</think>"
+                f"<tool name=\"calc\">{ans.split('=')[0].strip()}</tool>"
+                f"<result>{eval(ans.split('=')[0].strip())}</result>"
+                f"<think>That gives the first part; combine with the rest.</think> "
+                f"The answer is {ans}.</assistant><eos>")
+        out.append(conv)
+    return out
+
+
 # --------------------------------------------------------------------------
 # 3) RAG / retrieval examples (teach <tool name="retrieve"> and <context>)
 # --------------------------------------------------------------------------
@@ -149,7 +231,8 @@ def gen_rag(n=40000):
                     f"<result>{doc}</result>{a}</assistant><eos>")
         elif mode < 0.85:
             conv = (f"<bos><system>{RAG_SYSTEM}</system><user>{q}</user>"
-                    f"<assistant><context>{doc}</context>{a}</assistant><eos>")
+                    f"<assistant><think>Let me check the provided context.</think>"
+                    f"<context>{doc}</context>{a}</assistant><eos>")
         else:
             conv = (f"<bos><system>{RAG_SYSTEM}</system><user>{q}</user>"
                     f"<assistant><think>{doc}</think>{a}</assistant><eos>")
@@ -230,6 +313,7 @@ def main():
     ap.add_argument("--no-wiki", action="store_true")
     ap.add_argument("--no-rag", action="store_true")
     ap.add_argument("--no-glaive", action="store_true")
+    ap.add_argument("--no-mem", action="store_true")
     args = ap.parse_args()
 
     out_dir = args.out_dir
@@ -292,11 +376,19 @@ def main():
         except Exception as e:
             print(f"[prep] wikipedia skipped: {e}")
 
-    # --- synthetic tool data ---
+    # --- synthetic tool data (incl. interleaved reasoning) ---
     synth = gen_synthetic(int(60_000 * scale) + 60000)
-    n2 = pack(tok, synth, bin_path, bw_tool, SEQ_LEN)
+    inter = gen_interleaved(int(30_000 * scale) + 30000)
+    n2 = pack(tok, synth + inter, bin_path, bw_tool, SEQ_LEN)
     print(f"[prep] synthetic tool packed: {n2:,} tokens")
 
+    # --- learned-memory data ---
+    if not args.no_mem:
+        mem = gen_memory(int(30_000 * scale) + 30000)
+        nm = pack(tok, mem, bin_path, int(bw_tool * 0.6), SEQ_LEN)
+        print(f"[prep] memory packed: {nm:,} tokens")
+        n2 += nm
+
     # --- RAG / retrieval data ---
     if not args.no_rag:
         rag = gen_rag(int(40_000 * scale) + 40000)