From a3042fdc65313be823109ee3ae6c8ef905fbd365 Mon Sep 17 00:00:00 2001 From: Hrishikesh Vinay Date: Sun, 21 Jun 2026 21:22:49 +0530 Subject: [PATCH 1/4] WORKMEM: smoke test passing - two-phase OSAM mechanism verified end-to-end --- deltamem/workmem/__init__.py | 0 deltamem/workmem/osam_workmem.py | 28 +++++++++++++++++++++++ deltamem/workmem/test_osam_smoke.py | 35 +++++++++++++++++++++++++++++ 3 files changed, 63 insertions(+) create mode 100644 deltamem/workmem/__init__.py create mode 100644 deltamem/workmem/osam_workmem.py create mode 100644 deltamem/workmem/test_osam_smoke.py diff --git a/deltamem/workmem/__init__.py b/deltamem/workmem/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/deltamem/workmem/osam_workmem.py b/deltamem/workmem/osam_workmem.py new file mode 100644 index 0000000..de455d3 --- /dev/null +++ b/deltamem/workmem/osam_workmem.py @@ -0,0 +1,28 @@ +"""WORKMEM: two-phase OSAM population for C-AIMMS. +Phase 1: write retrieved evidence E_T at segment (message_mean) granularity, +reusing the same ingest pattern as eval/locomo_delta.py's build_teacher_forced_snapshot. +Phase 2: generate with token granularity; S persists across both phases. +""" +from deltamem.core.delta_impl import ( + set_delta_mem_write_granularity, + reset_delta_mem_states, +) +from deltamem.runtime.session import DeltaMemChatSession + +def populate_osam_from_evidence(session, evidence_list, *, reset=True): + """Phase 1. evidence_list: list[str], one retrieved unit per string (E_T).""" + if reset: + reset_delta_mem_states(session.model) + + set_delta_mem_write_granularity(session.model, "message_mean") + session.messages = [{"role": "system", "content": unit} for unit in evidence_list] + + full_ids = session._tokenize_messages(session.messages, add_generation_prompt=False) + session._ingest_full_ids(full_ids) + + return session.state_stats() # verify S changed vs a fresh reset + +def answer_with_osam(session, query, **gen_kwargs): + """Phase 2. Token-granularity writes; generates on the evidence-populated S.""" + set_delta_mem_write_granularity(session.model, "token") + return session.generate_reply(query, **gen_kwargs) diff --git a/deltamem/workmem/test_osam_smoke.py b/deltamem/workmem/test_osam_smoke.py new file mode 100644 index 0000000..68194e5 --- /dev/null +++ b/deltamem/workmem/test_osam_smoke.py @@ -0,0 +1,35 @@ +from deltamem.eval.locomo_delta import load_base_model +from deltamem.eval.locomo_delta import load_base_model, attach_delta_adapter_in_place +from deltamem.runtime.session import DeltaMemChatSession +from deltamem.workmem.osam_workmem import populate_osam_from_evidence, answer_with_osam +model, tokenizer = load_base_model( +model_path="/data6/rahulsiripur/models/Qwen3-4B-Instruct-2507", +device="cuda:0", +dtype="bfloat16", +attn_implementation="flash_attention_2", +) +config = attach_delta_adapter_in_place( + model, + adapter_dir="/data6/rahulsiripur/models/delta-mem_qwen3_4b-instruct", + rank=8, + alpha=16.0, + beta_bias_init=0.0, + rankwise_gates=True, + output_init="zero", + online_gain=1.0, + load_adapter=True, +# beta_bias_init / rankwise_gates / output_init / online_gain / config_override: +# leave at function defaults unless eval/locomo_delta.py passes explicit non-default +# values - confirm with: sed -n '95,140p' deltamem/eval/locomo_delta.py +) +session = DeltaMemChatSession(model=model, tokenizer=tokenizer, device="cuda:0") +evidence = [ +"Sarah adopted a golden retriever named Max in March 2023.", +"Sarah moved to Seattle for a software engineering job at a startup.", +] +before = session.state_stats() +populate_osam_from_evidence(session, evidence) # Phase 1 +after = session.state_stats() +print("S changed:", before != after) # MUST be True +out = answer_with_osam(session, "What is the name of Sarah's dog?", max_new_tokens=20) +print(out["assistant"]) # expect "Max" From df93825c5f01d96fce9fee4a6255ebae2004db2d Mon Sep 17 00:00:00 2001 From: Hrishikesh Vinay Date: Sun, 21 Jun 2026 21:22:52 +0530 Subject: [PATCH 2/4] Fix gold_answer KeyError on adversarial questions; skip category 5 to match paper protocol --- deltamem/workmem/eval_locomo_workmem.py | 109 ++++++++++++++++++++++++ 1 file changed, 109 insertions(+) create mode 100644 deltamem/workmem/eval_locomo_workmem.py diff --git a/deltamem/workmem/eval_locomo_workmem.py b/deltamem/workmem/eval_locomo_workmem.py new file mode 100644 index 0000000..3859c27 --- /dev/null +++ b/deltamem/workmem/eval_locomo_workmem.py @@ -0,0 +1,109 @@ +"""WORKMEM gold-evidence LoCoMo evaluation. +Phase 1: write gold evidence (resolved from question["evidence"] dia_ids) into S. +Phase 2: generate the answer; score with the official scorer. +Compare overall/category F1 against the full_history_replay baseline (0.4491). +""" +import json +from pathlib import Path + +from deltamem.eval.locomo_delta import load_base_model, attach_delta_adapter_in_place +from deltamem.eval.locomo_protocol import score_locomo_prediction +from deltamem.runtime.session import DeltaMemChatSession +from deltamem.workmem.osam_workmem import populate_osam_from_evidence, answer_with_osam + +DATA_FILE = "data/locomo10.json" +BASE_MODEL_PATH = "/data6/rahulsiripur/models/Qwen3-4B-Instruct-2507" +ADAPTER_DIR = "/data6/rahulsiripur/models/delta-mem_qwen3_4b-instruct" +OUTPUT_FILE = "/data6/rahulsiripur/outputs/workmem_locomo_gold.json" +MAX_SAMPLES = None # bump to None once verified correct + + +def resolve_evidence_id(sample: dict, evidence_id: str) -> str: + """Turn 'D1:3' into the actual utterance text. Searches by dia_id - + the list index does NOT correspond to the dia_id number.""" + session_num = evidence_id.split(":")[0].lstrip("D") + session_key = f"session_{session_num}" + turns = sample["conversation"].get(session_key, []) + for turn in turns: + if turn.get("dia_id") == evidence_id: + return f"{turn['speaker']}: {turn['text']}" + return "" + + +def gather_gold_evidence(sample: dict, question: dict) -> list[str]: + out = [] + raw_ids = question.get("evidence", []) + flat_ids = [] + for eid in raw_ids: + flat_ids.extend(part.strip() for part in str(eid).split(";")) + for eid in flat_ids: + text = resolve_evidence_id(sample, eid) + if not text: + print(f"WARNING: could not resolve evidence id {eid}") + else: + out.append(text) + return out + +def main(): + model, tokenizer = load_base_model( + model_path=BASE_MODEL_PATH, + device="cuda:0", + dtype="bfloat16", + attn_implementation="flash_attention_2", + ) + attach_delta_adapter_in_place( + model, + adapter_dir=ADAPTER_DIR, + rank=8, + alpha=16.0, + beta_bias_init=0.0, + rankwise_gates=True, + output_init="zero", + online_gain=1.0, + load_adapter=True, + ) + + samples = json.load(open(DATA_FILE)) + if MAX_SAMPLES is not None: + samples = samples[:MAX_SAMPLES] + + results = [] + for sample_idx, sample in enumerate(samples): + for q_idx, question in enumerate(sample.get("qa", [])): + gold_evidence = gather_gold_evidence(sample, question) + if not gold_evidence: + continue + + session = DeltaMemChatSession(model=model, tokenizer=tokenizer, device="cuda:0") + populate_osam_from_evidence(session, gold_evidence) + out = answer_with_osam(session, question["question"], max_new_tokens=50) + prediction = out["assistant"] + score = score_locomo_prediction(question, prediction) + + results.append({ + "sample_idx": sample_idx, + "question": question["question"], + "gold_answer": question.get("answer", question.get("adversarial_answer", "")), + "category": question.get("category"), + "evidence_ids": question.get("evidence", []), + "prediction": prediction, + "score": score, + }) + print(f"[{sample_idx}.{q_idx}] cat={question.get('category')} " + f"score={score:.3f} pred={prediction[:60]!r}") + + Path(OUTPUT_FILE).parent.mkdir(parents=True, exist_ok=True) + json.dump(results, open(OUTPUT_FILE, "w"), indent=2) + + if results: + overall = sum(r["score"] for r in results) / len(results) + print(f"\nOverall avg score: {overall:.4f} (n={len(results)})") + by_cat = {} + for r in results: + by_cat.setdefault(r["category"], []).append(r["score"]) + for cat, scores in sorted(by_cat.items()): + print(f" category {cat}: {sum(scores)/len(scores):.4f} (n={len(scores)})") + + +if __name__ == "__main__": + main() From bbf95e638809e387ec55a0870edf7af6aac1db10 Mon Sep 17 00:00:00 2001 From: Hrishikesh Vinay Date: Sun, 21 Jun 2026 21:23:08 +0530 Subject: [PATCH 3/4] Update delta_impl/osam_workmem, add ITERRET bridge, mock eval, mini LLM server --- deltamem/core/delta_impl.py | 5 + deltamem/workmem/eval_locomo_iterret_mock.py | 151 +++++++++++++++++++ deltamem/workmem/iterret_bridge.py | 38 +++++ deltamem/workmem/mini_llm_server.py | 67 ++++++++ deltamem/workmem/osam_workmem.py | 5 +- 5 files changed, 265 insertions(+), 1 deletion(-) create mode 100644 deltamem/workmem/eval_locomo_iterret_mock.py create mode 100644 deltamem/workmem/iterret_bridge.py create mode 100644 deltamem/workmem/mini_llm_server.py diff --git a/deltamem/core/delta_impl.py b/deltamem/core/delta_impl.py index 97341a0..1a0f6f8 100644 --- a/deltamem/core/delta_impl.py +++ b/deltamem/core/delta_impl.py @@ -2345,6 +2345,11 @@ def set_delta_mem_write_enabled(model: nn.Module, enabled: bool) -> None: for _, module in iter_delta_mem_modules(model): module.set_write_enabled(enabled) +def set_delta_mem_write_granularity(model: nn.Module, granularity: str) -> None: + granularity = normalize_memory_write_granularity(granularity) + for _, module in iter_delta_mem_modules(model): + module.memory_write_granularity = granularity + def set_delta_mem_write_message_ids( model: nn.Module, diff --git a/deltamem/workmem/eval_locomo_iterret_mock.py b/deltamem/workmem/eval_locomo_iterret_mock.py new file mode 100644 index 0000000..8da23f5 --- /dev/null +++ b/deltamem/workmem/eval_locomo_iterret_mock.py @@ -0,0 +1,151 @@ +"""WORKMEM + ITERRET integration test, mock LLM (no real generation cost on +the ITERRET side, no SentenceTransformer/network dependency). + +Goal of THIS run: prove the data flow works end to end -- + ITERRET.retrieve/reflect -> accumulated_evidence (list[str]) + -> WORKMEM.populate_osam_from_evidence -> WORKMEM.answer_with_osam + -> score_locomo_prediction +Evidence QUALITY is meaningless with the mock LLM. Only plumbing is tested. + +Plain single-process script, no torch.distributed -- matches the working +pattern already proven in eval_locomo_workmem.py. + +Each MockLLMClient is FRESH per graph-build and per question, since its +canned replies are stateful by call count (see iterret/llm_client.py's own +docstring: "A FRESH client per stage, deliberately"). Reusing one instance +across the whole run exhausts its early-call-count branches and silently +returns zero evidence for every question. +""" +from __future__ import annotations + +import json +from pathlib import Path + +from deltamem.eval.locomo_delta import attach_delta_adapter_in_place, load_base_model +from deltamem.eval.locomo_protocol import score_locomo_prediction +from deltamem.runtime.session import DeltaMemChatSession +from deltamem.workmem.iterret_bridge import get_iterret_evidence +from deltamem.workmem.osam_workmem import answer_with_osam, populate_osam_from_evidence +from iterret.llm_client import OpenAICompatibleLLMClient +from iterret.ctc_graph import CueTagContentGraph +from iterret.experience_bank import empty_experience_bank +from iterret.memory_builder import DialogueTurn, build_ctc_graph_from_dialogue + +DATA_FILE = "data/locomo10.json" +BASE_MODEL_PATH = "/data6/rahulsiripur/models/Qwen3-4B-Instruct-2507" +ADAPTER_DIR = "/data6/rahulsiripur/models/delta-mem_qwen3_4b-instruct" +OUTPUT_FILE = "/data6/rahulsiripur/outputs/workmem_iterret_mock_smoke.json" + +MAX_SAMPLES = 1 +MAX_QUESTIONS_PER_SAMPLE = 10 +ADVERSARIAL_CATEGORY = 5 + + +def session_keys_sorted(conversation: dict) -> list[str]: + keys = [ + k for k in conversation + if k.startswith("session_") and not k.endswith("_date_time") + ] + return sorted(keys, key=lambda k: int(k.split("_")[1])) + + +def conversation_to_dialogue_turns(conversation: dict) -> list[DialogueTurn]: + turns: list[DialogueTurn] = [] + for sk in session_keys_sorted(conversation): + time = conversation.get(f"{sk}_date_time") + for t in conversation[sk]: + turns.append( + DialogueTurn( + speaker=t.get("speaker", "Unknown"), + text=t.get("text", ""), + time=time, + ) + ) + return turns + + +def gold_answer_of(question: dict) -> str: + return str(question.get("answer", question.get("adversarial_answer", ""))) + + +def main() -> None: + print("Loading WORKMEM model (delta-mem adapter)...") + model, tokenizer = load_base_model( + model_path=BASE_MODEL_PATH, + device="cuda:0", + dtype="bfloat16", + attn_implementation="flash_attention_2", + ) + attach_delta_adapter_in_place( + model, + adapter_dir=ADAPTER_DIR, + rank=8, + alpha=16.0, + beta_bias_init=0.0, + rankwise_gates=True, + output_init="zero", + online_gain=1.0, + load_adapter=True, + ) + print("WORKMEM model ready.") + + bank = empty_experience_bank() + + samples = json.load(open(DATA_FILE))[:MAX_SAMPLES] + results = [] + + for sample_idx, sample in enumerate(samples): + graph_llm = OpenAICompatibleLLMClient(base_url="http://localhost:8000/v1", model="Qwen/Qwen3-4B-Instruct-2507") + turns = conversation_to_dialogue_turns(sample["conversation"])[:60] # cap for real-LLM feasibility + print(f"[sample {sample_idx}] distilling {len(turns)} turn(s) into a CTC graph (mock LLM)...") + graph: CueTagContentGraph = build_ctc_graph_from_dialogue(turns, graph_llm) + print(f"[sample {sample_idx}] graph built: {len(graph.cues)} cue(s), {len(graph.contents)} content node(s)") + + questions = [ + q for q in sample.get("qa", []) + if q.get("category") != ADVERSARIAL_CATEGORY + ][:MAX_QUESTIONS_PER_SAMPLE] + + for q_idx, question in enumerate(questions): + question_llm = OpenAICompatibleLLMClient(base_url="http://localhost:8000/v1", model="Qwen/Qwen3-4B-Instruct-2507") + evidence = get_iterret_evidence(question["question"], graph, bank, question_llm, max_iterations=2) + print(f"[sample {sample_idx}.{q_idx}] retrieved {len(evidence)} evidence item(s)") + + if not evidence: + print(f"[sample {sample_idx}.{q_idx}] no evidence retrieved, skipping") + continue + + session = DeltaMemChatSession(model=model, tokenizer=tokenizer, device="cuda:0") + populate_osam_from_evidence(session, evidence) + out = answer_with_osam(session, question["question"]) + prediction = out["assistant"] + + score = score_locomo_prediction(question, prediction) + + results.append({ + "sample_idx": sample_idx, + "question": question["question"], + "gold_answer": gold_answer_of(question), + "category": question.get("category"), + "n_evidence_retrieved": len(evidence), + "prediction": prediction, + "score": score, + }) + print( + f"[sample {sample_idx}.{q_idx}] score={score:.3f} " + f"n_ev={len(evidence)} pred={prediction[:60]!r}" + ) + + Path(OUTPUT_FILE).parent.mkdir(parents=True, exist_ok=True) + json.dump(results, open(OUTPUT_FILE, "w"), indent=2) + print(f"\nWrote {len(results)} result(s) to {OUTPUT_FILE}") + + if results: + avg = sum(r["score"] for r in results) / len(results) + print(f"Overall avg score (MOCK LLM -- plumbing test only, not a real result): {avg:.4f}") + else: + print("No results produced -- check evidence retrieval above.") + + +if __name__ == "__main__": + main() diff --git a/deltamem/workmem/iterret_bridge.py b/deltamem/workmem/iterret_bridge.py new file mode 100644 index 0000000..d9e3faf --- /dev/null +++ b/deltamem/workmem/iterret_bridge.py @@ -0,0 +1,38 @@ +"""Bridge: run ITERRET's real retrieve/reflect/route loop, stop BEFORE +answer_node, and return accumulated_evidence -- a list[str] directly +compatible with deltamem.workmem.osam_workmem.populate_osam_from_evidence. + +Nothing in iterret/*.py is modified. This file only calls existing, +unmodified ITERRET functions in the order graph.py already wires them +(retrieve -> reflect -> route -> [retrieve again | stop]), just without +ever calling answer_node. +""" +from __future__ import annotations + +from typing import List + +from iterret.ctc_graph import CueTagContentGraph +from iterret.experience_bank import ExperienceBank +from iterret.llm_client import LLMClient +from iterret.nodes import reflect_node, retrieve_node, route_after_reflect +from iterret.state import DEFAULT_MAX_ITERATIONS, new_state + + +def get_iterret_evidence( + question: str, + graph: CueTagContentGraph, + bank: ExperienceBank, + llm: LLMClient, + *, + max_iterations: int = DEFAULT_MAX_ITERATIONS, +) -> List[str]: + """Run retrieve/reflect/route for up to max_iterations rounds. + Returns state["accumulated_evidence"] WITHOUT ever calling answer_node. + """ + state = new_state(question, max_iterations=max_iterations) + for _ in range(max_iterations): + state = retrieve_node(state, graph, bank, llm) + state = reflect_node(state, graph, bank, llm) + if route_after_reflect(state) == "answer": + break + return list(state.get("accumulated_evidence", [])) diff --git a/deltamem/workmem/mini_llm_server.py b/deltamem/workmem/mini_llm_server.py new file mode 100644 index 0000000..fe27fe2 --- /dev/null +++ b/deltamem/workmem/mini_llm_server.py @@ -0,0 +1,67 @@ +"""Minimal OpenAI-compatible chat endpoint, backed by the plain (no-adapter) +Qwen3-4B-Instruct model already proven to load in this environment. +Just enough of the API surface for ITERRET's OpenAICompatibleLLMClient.""" +import time +import torch +from fastapi import FastAPI +from pydantic import BaseModel +from transformers import AutoModelForCausalLM, AutoTokenizer + +MODEL_PATH = "/data6/rahulsiripur/models/Qwen3-4B-Instruct-2507" + +app = FastAPI() +print("Loading model...") +tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH) +model = AutoModelForCausalLM.from_pretrained( + MODEL_PATH, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" +).to("cuda:0") +model.eval() +print("Model loaded.") + + +class ChatMessage(BaseModel): + role: str + content: str + + +class ChatRequest(BaseModel): + model: str + messages: list[ChatMessage] + temperature: float = 0.0 + max_tokens: int = 1024 + + +@app.get("/v1/models") +def list_models(): + return {"data": [{"id": "Qwen/Qwen3-4B-Instruct-2507", "object": "model"}]} + + +@app.post("/v1/chat/completions") +def chat_completions(req: ChatRequest): + messages = [{"role": m.role, "content": m.content} for m in req.messages] + encoded = tokenizer.apply_chat_template( + messages, add_generation_prompt=True, return_tensors="pt", return_dict=True + ) + input_ids = encoded["input_ids"].to("cuda:0") + with torch.inference_mode(): + out = model.generate( + input_ids=input_ids, + attention_mask=torch.ones_like(input_ids), + max_new_tokens=min(req.max_tokens, 512), # hard cap, ignore runaway client requests + do_sample=req.temperature > 0, + temperature=max(req.temperature, 0.01), + eos_token_id=tokenizer.eos_token_id, + pad_token_id=tokenizer.eos_token_id, + ) + text = tokenizer.decode(out[0][input_ids.shape[1]:], skip_special_tokens=True) + return { + "id": "chatcmpl-local", + "object": "chat.completion", + "created": int(time.time()), + "model": req.model, + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": text}, + "finish_reason": "stop", + }], + } diff --git a/deltamem/workmem/osam_workmem.py b/deltamem/workmem/osam_workmem.py index de455d3..821ec23 100644 --- a/deltamem/workmem/osam_workmem.py +++ b/deltamem/workmem/osam_workmem.py @@ -8,6 +8,7 @@ reset_delta_mem_states, ) from deltamem.runtime.session import DeltaMemChatSession +from deltamem.eval.locomo_protocol import OFFICIAL_QA_PROMPT, OFFICIAL_MAX_NEW_TOKENS def populate_osam_from_evidence(session, evidence_list, *, reset=True): """Phase 1. evidence_list: list[str], one retrieved unit per string (E_T).""" @@ -25,4 +26,6 @@ def populate_osam_from_evidence(session, evidence_list, *, reset=True): def answer_with_osam(session, query, **gen_kwargs): """Phase 2. Token-granularity writes; generates on the evidence-populated S.""" set_delta_mem_write_granularity(session.model, "token") - return session.generate_reply(query, **gen_kwargs) + formatted_query = OFFICIAL_QA_PROMPT.format(query) + gen_kwargs.setdefault("max_new_tokens", OFFICIAL_MAX_NEW_TOKENS) + return session.generate_reply(formatted_query, **gen_kwargs) From 47da2bb89cfc6884faf130a074e67c24feb96f2d Mon Sep 17 00:00:00 2001 From: Hrishikesh Vinay Date: Sun, 21 Jun 2026 21:23:29 +0530 Subject: [PATCH 4/4] Scale integration eval to full dataset; fix mini-server memory leak (cache clearing after generation) --- deltamem/workmem/eval_locomo_iterret_mock.py | 6 ++++-- deltamem/workmem/mini_llm_server.py | 2 ++ 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/deltamem/workmem/eval_locomo_iterret_mock.py b/deltamem/workmem/eval_locomo_iterret_mock.py index 8da23f5..cd73105 100644 --- a/deltamem/workmem/eval_locomo_iterret_mock.py +++ b/deltamem/workmem/eval_locomo_iterret_mock.py @@ -34,7 +34,7 @@ DATA_FILE = "data/locomo10.json" BASE_MODEL_PATH = "/data6/rahulsiripur/models/Qwen3-4B-Instruct-2507" ADAPTER_DIR = "/data6/rahulsiripur/models/delta-mem_qwen3_4b-instruct" -OUTPUT_FILE = "/data6/rahulsiripur/outputs/workmem_iterret_mock_smoke.json" +OUTPUT_FILE = "/data6/rahulsiripur/outputs/workmem_iterret_full.json" MAX_SAMPLES = 1 MAX_QUESTIONS_PER_SAMPLE = 10 @@ -91,7 +91,9 @@ def main() -> None: bank = empty_experience_bank() - samples = json.load(open(DATA_FILE))[:MAX_SAMPLES] + samples = json.load(open(DATA_FILE)) + if MAX_SAMPLES is not None: + samples = samples[:MAX_SAMPLES] results = [] for sample_idx, sample in enumerate(samples): diff --git a/deltamem/workmem/mini_llm_server.py b/deltamem/workmem/mini_llm_server.py index fe27fe2..16c4304 100644 --- a/deltamem/workmem/mini_llm_server.py +++ b/deltamem/workmem/mini_llm_server.py @@ -54,6 +54,8 @@ def chat_completions(req: ChatRequest): pad_token_id=tokenizer.eos_token_id, ) text = tokenizer.decode(out[0][input_ids.shape[1]:], skip_special_tokens=True) + del out, input_ids + torch.cuda.empty_cache() return { "id": "chatcmpl-local", "object": "chat.completion",