docs: add Generalist model card and fix dtype metadata

README.md

@@ -2,24 +2,101 @@
 license: apache-2.0
 language:
 - en
+library_name: transformers
+pipeline_tag: text-generation
 base_model:
 - Qwen/Qwen3-4B-Instruct-2507
-library_name: transformers
+base_model_relation: finetune
+tags:
+- memory-agent
+- reinforcement-learning
+- long-context
+- tool-use
+- qwen3
+- grpo
+arxiv: 2602.18493
 ---
 
-# UMA-4B
+# UMA-4B (Generalist)
+
+**UMA-4B** is the Generalist checkpoint of the Unified Memory Agent (UMA) introduced in [Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning](https://arxiv.org/abs/2602.18493).
+
+UMA is a tool-using memory agent that incrementally maintains a compact core summary and a structured key-value Memory Bank. The same policy performs memory construction and downstream question answering through explicit memory and retrieval operations.
+
+- **Code:** [github.com/ictnlp/unified-memory-agent](https://github.com/ictnlp/unified-memory-agent)
+- **Paper:** [arXiv:2602.18493](https://arxiv.org/abs/2602.18493)
+- **Specialist checkpoint:** [ICTNLP/UMA-LedgerQA-4B](https://huggingface.co/ICTNLP/UMA-LedgerQA-4B)
+
+## Checkpoint Variant
+
+This repository contains the **Generalist** UMA checkpoint used for the paper's Test-Time Learning and Accurate Retrieval evaluations.
+
+| Property | Value |
+| --- | --- |
+| Base model | `Qwen/Qwen3-4B-Instruct-2507` |
+| Parameters | 4B |
+| Weight format | BF16 Safetensors |
+| Training method | End-to-end reinforcement learning with Task-Stratified GRPO |
+| Training data | HotpotQA and the Mem-alpha corpus |
+| Ledger-QA training data | None |
+| Reported default context budget | 16K |
+
+The Generalist and Specialist checkpoints share the same UMA architecture and tool interface. The Specialist checkpoint is additionally adapted to Ledger-QA; use this Generalist checkpoint for the broader cross-task setting.
+
+## Intended Use
+
+This checkpoint is intended for research on:
 
-Agentic RL fine-tuned model
+- long-context and streaming memory agents;
+- proactive structured memory construction;
+- memory maintenance with explicit tool calls;
+- downstream question answering over reusable memory;
+- evaluation and extension of the UMA framework.
 
-## Usage
+The checkpoint is designed to run inside the UMA two-phase agent loop. A plain text-generation call loads the language model, but does not by itself instantiate the Memory Bank, retrieval tools, prompts, or memory-to-QA workflow.
+
+## Loading the Weights
 
 ```python
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "ICTNLP/UMA-4B"
+
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+```
+
+The model can also be served through an OpenAI-compatible inference server:
 
-tokenizer = AutoTokenizer.from_pretrained("dp66/UMA-4B")
-model = AutoModelForCausalLM.from_pretrained("dp66/UMA-4B")
+```bash
+vllm serve ICTNLP/UMA-4B \
+  --max-model-len 16384 \
+  --gpu-memory-utilization 0.8
 ```
 
-## Training Details
+For full memory-agent inference, including the Memory Bank, memory tools, embedding retrieval, prompts, and benchmark runners, follow the [official repository](https://github.com/ictnlp/unified-memory-agent).
+
+## Limitations
+
+- The checkpoint is a research model and may generate incorrect answers or perform incorrect memory updates.
+- Agent behavior depends on the UMA prompt templates, tool implementations, retrieval backend, chunking policy, and inference configuration.
+- The model was primarily trained and evaluated on English-language research benchmarks.
+- Persistent-memory applications can involve sensitive information. Deployments should provide appropriate privacy controls, retention policies, and user oversight.
+- This checkpoint should not be used as the sole basis for high-stakes decisions.
+
+## Citation
+
+```bibtex
+@article{zhang2026learning,
+  title   = {Learning to Remember: End-to-End Training of Memory Agents for Long-Context Reasoning},
+  author  = {Zhang, Kehao and Gui, Shangtong and Yang, Sheng and Chen, Wei and Feng, Yang},
+  journal = {arXiv preprint arXiv:2602.18493},
+  year    = {2026}
+}
+```
 
-- Base Model: Qwen/Qwen3-4B-Instruct-2507

config.json

@@ -4,7 +4,7 @@
   ],
   "attention_bias": false,
   "attention_dropout": 0.0,
-  "dtype": "float32",
+  "dtype": "bfloat16",
   "eos_token_id": 151645,
   "head_dim": 128,
   "hidden_act": "silu",