Workmem osam

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,

deltamem/workmem/eval_locomo_iterret_mock.py

@@ -0,0 +1,153 @@
+"""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_full.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))
+    if MAX_SAMPLES is not None:
+        samples = samples[: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()

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()

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", []))

deltamem/workmem/mini_llm_server.py

@@ -0,0 +1,69 @@
+"""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)
+    del out, input_ids
+    torch.cuda.empty_cache()
+    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",
+        }],
+    }