Update README.md

README.md

@@ -1,3 +1,227 @@
----
-license: apache-2.0
----
+---
+license: apache-2.0
+language:
+- en
+- zh
+base_model:
+- Qwen/Qwen3-Embedding-4B
+tags:
+- embedding
+- retriever
+- RAG
+---
+
+# Mindscape-Aware RAG (MiA-RAG)
+
+[![Paper](https://img.shields.io/badge/Paper-arXiv%3A2512.17220-red)](https://arxiv.org/pdf/2512.17220)
+[![Model](https://img.shields.io/badge/HuggingFace-MiA--Emb--4B-yellow)](https://huggingface.co/MindscapeRAG/MiA-Emb-4B)
+
+This repository provides the inference implementation for **MiA-Emb (Mindscape-Aware Embedding)**, the retriever component in the **MiA-RAG** framework.
+
+**MiA-RAG** introduces explicit **global context awareness** via a **Mindscape**—a document-level semantic scaffold constructed by **hierarchical summarization**. By conditioning **both retrieval and generation** on the same Mindscape, MiA-RAG enables globally grounded retrieval and more coherent long-context reasoning.
+
+---
+
+## ✨ Key Features
+
+- **Mindscape as Global Semantic Scaffold**  
+  Builds a Mindscape through **hierarchical bottom-up summarization** (chunk summaries → global summary) and uses it as persistent global memory.
+
+- **Mindscape-Aware Capabilities**  
+  Supports the three core benefits for long-context understanding:
+  - **Enriched Understanding**: fill in missing context and resolve underspecified meanings
+  - **Selective Retrieval**: bias retrieval toward the active topic’s semantic frame
+  - **Integrative Reasoning**: interpret retrieved evidence within a coherent global context
+
+- **Dual-Granularity Retrieval**  
+  - **Chunk Retrieval** for narrative passages (standard RAG)
+  - **Node Retrieval** for knowledge graph entities (GraphRAG-style)
+
+- **State-of-the-Art Retrieval Performance**  
+  Strong results on long-context benchmarks such as NarrativeQA and DetectiveQA, outperforming strong baselines including Qwen3-Embedding and [SitEmb](https://huggingface.co/SituatedEmbedding/SitEmb-v1.5-Qwen3).
+
+
+---
+
+
+## 🚀 Usage
+
+### Installation
+
+```bash
+pip install torch transformers>=4.53.0
+```
+
+---
+
+### 1) Initialization
+
+> MiA-Emb-4B is initialized from **`Qwen3-Embedding-4B`**.
+
+```python
+import torch
+import torch.nn.functional as F
+from transformers import AutoTokenizer, AutoModel
+
+# Configuration
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+# Inference Parameters
+residual = True         # Enable residual connection logic
+residual_factor = 0.5   # Balance between local and global
+node_delimiter = "<|repo_name|>"  # Special token for Node tasks
+
+# Load Tokenizer (base)
+tokenizer = AutoTokenizer.from_pretrained(
+    "Qwen/Qwen3-Embedding-4B",
+    trust_remote_code=True,
+    padding_side="left"
+)
+
+# Load Model
+model = AutoModel.from_pretrained(
+    "MindscapeRAG/MiA-Emb-4B",
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16,
+    attn_implementation="flash_attention_2",
+    device_map={"": 0}
+)
+```
+
+---
+
+### 2) Chunk Retrieval
+
+Use this mode to retrieve narrative text chunks. A **Global Summary** is injected into the prompt as the “Mindscape”.
+
+```python
+def get_query_prompt(query, summary="", residual=False):
+    """Construct input prompt with global summary (Eq. 5 in paper)."""
+    task_desc = "Given a search query with the book's summary, retrieve relevant chunks or helpful entities summaries from the given context that answer the query"
+    summary_prefix = "\n\nHere is the summary providing possibly useful global information. Please encode the query based on the summary:\n"
+
+    # Insert PAD token to capture residual embedding before the summary
+    middle_token = tokenizer.pad_token if residual else ""
+
+    return (
+        f"Instruct: {task_desc}\n"
+        f"Query: {query}{middle_token}{summary_prefix}{summary}{node_delimiter}"
+    )
+
+def encode_chunk(texts, is_query=False, residual=False):
+    batch = tokenizer(
+        texts,
+        max_length=4096,
+        padding=True,
+        truncation=True,
+        return_tensors="pt"
+    ).to(model.device)
+
+    outputs = model(**batch)
+
+    # 1) Main Embedding (Last Token)
+    emb_main = last_token_pool(outputs.last_hidden_state, batch["attention_mask"])
+
+    # 2) Residual Embedding (PAD Token)
+    emb_res = None
+    if residual and is_query:
+        emb_res = extract_residual_token(outputs, batch, tokenizer.pad_token_id)
+
+    emb_main = F.normalize(emb_main, p=2, dim=-1)
+    emb_res = F.normalize(emb_res, p=2, dim=-1) if emb_res is not None else None
+    return emb_main, emb_res
+
+
+# --- Example ---
+query = "Who is the protagonist?"
+global_summ = "A summary of the entire book..."
+chunk = "Harry looked at the scar on his forehead."
+
+# Encode
+q_emb, q_res = encode_chunk(
+    [get_query_prompt(query, global_summ, residual=True)],
+    is_query=True,
+    residual=True
+)
+c_emb, _ = encode_chunk([chunk], is_query=False)
+
+# Score Fusion
+score = q_emb @ c_emb.T
+if q_res is not None:
+    score = (1 - residual_factor) * score + residual_factor * (q_res @ c_emb.T)
+
+print(f"Chunk Similarity: {score.item():.4f}")
+```
+
+---
+
+### 3) Node Retrieval
+
+MiA-Emb can retrieve knowledge graph entities (**Nodes**). This mode extracts embeddings from the `<|repo_name|>` token position.
+
+**Candidate format:**
+`Entity Name : Entity Description`
+
+Example:
+`Mary Campbell Smith : Mary Campbell Smith is mentioned as the translator...`
+
+```python
+def encode_node_query(texts, residual=True, node_delimiter="<|repo_name|>"):
+    batch = tokenizer(texts, padding=True, return_tensors="pt").to(model.device)
+    outputs = model(**batch)
+
+    # 1) Node Main Embedding: extract from <|repo_name|> position
+    node_id = tokenizer.encode(node_delimiter, add_special_tokens=False)[0]
+    q_emb_node = extract_specific_token(outputs, batch, node_id)
+
+    # 2) Residual Embedding: extract from [PAD] position
+    q_emb_res = extract_residual_token(outputs, batch, tokenizer.pad_token_id) if residual else None
+
+    q_emb_node = F.normalize(q_emb_node, p=2, dim=-1)
+    q_emb_res = F.normalize(q_emb_res, p=2, dim=-1) if q_emb_res is not None else None
+    return q_emb_node, q_emb_res
+
+
+# --- Example ---
+query = "Who is the protagonist?"
+global_summ = "A summary of the entire book..."
+
+# 1) Encode Query (Node Token)
+q_emb_node, q_emb_res = encode_node_query(
+    [get_query_prompt(query, global_summ, residual=True)],
+    residual=True
+)
+
+# 2) Encode Entity Candidate
+entity_text = "Harry Potter : The main protagonist of the series..."
+n_emb, _ = encode_chunk([entity_text], is_query=False)
+
+# 3) Score Fusion
+final_score = (1 - residual_factor) * (q_emb_node @ n_emb.T)
+if q_emb_res is not None:
+    final_score = final_score + residual_factor * (q_emb_res @ n_emb.T)
+
+print(f"Node Similarity: {final_score.item():.4f}")
+```
+
+---
+
+
+## 📜 Citation
+
+
+
+If you find this work useful, please cite:
+
+```bibtex
+@misc{li2025mindscapeawareretrievalaugmentedgeneration,
+      title={Mindscape-Aware Retrieval Augmented Generation for Improved Long Context Understanding}, 
+      author={Yuqing Li and Jiangnan Li and Zheng Lin and Ziyan Zhou and Junjie Wu and Weiping Wang and Jie Zhou and Mo Yu},
+      year={2025},
+      eprint={2512.17220},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL},
+      url={https://arxiv.org/abs/2512.17220}, 
+}
+```