Add pipeline tag, paper link, and improve documentation

README.md

@@ -1,9 +1,15 @@
 ---
 license: apache-2.0
+pipeline_tag: text-retrieval
 ---
+
 # EmbeddingRWKV
 
-A high-efficiency text embedding and reranking model based on RWKV architecture.
+EmbeddingRWKV is a high-efficiency text embedding and reranking model based on the RWKV architecture, introduced in the paper [EmbeddingRWKV: State-Centric Retrieval with Reusable States](https://huggingface.co/papers/2601.07861).
+
+It utilizes **State-Centric Retrieval**, a unified retrieval paradigm that uses "states" as a bridge to connect embedding models and rerankers, significantly improving inference speed for reranking tasks.
+
+[**Paper**](https://huggingface.co/papers/2601.07861) | [**GitHub**](https://github.com/howard-hou/EmbeddingRWKV)
 
 ## 📦 Installation
 
@@ -48,7 +54,7 @@ documents = [
 ]
 # Encode Query
 q_tokens = tokenizer.encode(query, add_eos=True)
-q_emb, _ = emb_model.forward(q_tokens, None) # shape: [1, Dim]
+q_emb, _ = emb_model.forward_text_only(q_tokens, None) # shape: [1, Dim]
 
 # Encode Documents (Batch)
 doc_batch = [tokenizer.encode(doc, add_eos=True) for doc in documents]
@@ -58,17 +64,16 @@ for i in range(len(doc_batch)):
     # Prepend 0s (Left Padding)
     doc_batch[i] = [0] * pad_len + doc_batch[i]
 
-d_embs, _ = emb_model.forward(doc_batch, None)
+d_embs, _ = emb_model.forward_text_only(doc_batch, None)
 
 # Calculate Cosine Similarity
 scores_emb = F.cosine_similarity(q_emb, d_embs)
-print("\nEmbeddingRWKV Cosine Similarity:")
+print("
+EmbeddingRWKV Cosine Similarity:")
 for doc, score in zip(documents, scores_emb):
     print(f"[{score.item():.4f}] {doc}")
 ```
 
-For production use cases, running inference in batches is significantly faster.
-
 ### ⚠️ Critical Performance Tip: Pad to Same Length
 
 While the model supports batches with variable sequence lengths, **we strongly recommend padding all sequences to the same length** for maximum GPU throughput.
@@ -81,14 +86,7 @@ While the model supports batches with variable sequence lengths, **we strongly r
 
 The `RWKVReRanker` utilizes the final hidden state produced by the main `EmbeddingRWKV` model to score the relevance between a query and a document.
 
-### Online Mode
-
-#### Workflow
-1.  **Format** Query and Document based on Online template.
-2.  Run the **Embedding Model** to generate the final State.
-3.  Feed the **TimeMixing State** (`state[1]`) into the **ReRanker** to get a relevance score.
-
-#### 📝 Online Mode Usage Example
+### Online Mode Usage Example
 
 ```python
 import torch
@@ -96,7 +94,6 @@ from rwkv_emb.tokenizer import RWKVTokenizer
 from rwkv_emb.model import EmbeddingRWKV, RWKVReRanker
 
 # 1. Load Models
-# The ReRanker weights are stored in the differernt checkpoint
 emb_model = EmbeddingRWKV(model_path='/path/to/EmbeddingRWKV.pth')
 reranker = RWKVReRanker(model_path='/path/to/RWKVReRanker.pth')
 
@@ -111,34 +108,30 @@ documents = [
 ]
 
 # 3. Construct Input Pairs
-# We treat the Query and Document as a single sequence.
 pairs = []
-online_template = "Instruct: Given a query, retrieve documents that answer the query\nDocument: {document}\nQuery: {query}"
+online_template = "Instruct: Given a query, retrieve documents that answer the query
+Document: {document}
+Query: {query}"
 for doc in documents:
-    # Format: Instruct + Document + Query
     text = online_template.format(document=doc, query=query)
     pairs.append(text)
 
-# 4. Tokenize & Pad (Critical for Batch Performance)
+# 4. Tokenize & Pad
 batch_tokens = [tokenizer.encode(p, add_eos=True) for p in pairs]
-
-# Left pad to same length for efficiency
 max_len = max(len(t) for t in batch_tokens)
 for i in range(len(batch_tokens)):
     batch_tokens[i] = [0] * (max_len - len(batch_tokens[i])) + batch_tokens[i]
 
 # 5. Get States from Embedding Model
-# We don't need the embedding output here, we only need the final 'state'
 _, state = emb_model.forward(batch_tokens, None)
 
 # 6. Score with ReRanker
-# The ReRanker expects the TimeMixing State: state[1]
-# state[1] shape: [Layers, Batch, Heads, HeadSize, HeadSize]
 logits = reranker.forward(state[1])
-scores = torch.sigmoid(logits) # Convert logits to probabilities (0-1)
+scores = torch.sigmoid(logits) 
 
 # 7. Print Results
-print("\nRWKVReRanker Online Scores:")
+print("
+RWKVReRanker Online Scores:")
 for doc, score in zip(documents, scores):
     print(f"[{score:.4f}] {doc}")
 ```
@@ -146,71 +139,39 @@ for doc, score in zip(documents, scores):
 ### Offline Mode (Cached Doc State)
 For scenarios where documents are static but queries change (e.g., Search Engines, RAG), you can **pre-compute and cache the document states**. This reduces query-time latency from O(L_doc + L_query) to just O(L_query).
 
-#### Workflow
-
-1.  **Indexing**: Process `Instruct + Document` -\> Save State.
-2.  **Querying**: Load State -\> Process `Query` -\> Score.
-
-#### 📝 Offline Mode Usage Example
-
 ```python
 # --- Phase 1: Indexing (Pre-computation) ---
-# Note: Do NOT add EOS here, because the sequence continues with the query later.
-doc_template = "Instruct: Given a query, retrieve documents that answer the query\nDocument: {document}\n"
+doc_template = "Instruct: Given a query, retrieve documents that answer the query
+Document: {document}
+"
 cached_states = []
 
-print("Indexing documents...")
 for doc in documents:
     text = doc_template.format(document=doc)
-    # add_eos=False is CRITICAL here
     tokens = tokenizer.encode(text, add_eos=False) 
-    
-    # Forward pass
     _, state = emb_model.forward(tokens, None)
-    
-    # Move state to CPU to save GPU memory during storage
-    # State structure: [Tensor(Tokenshift), Tensor(TimeMix)]
     cpu_state = [s.cpu() for s in state]
     cached_states.append(cpu_state)
-# Save cached states to disk (optional)
-torch.save(cached_states, 'cached_doc_states.pth')
 
 # --- Phase 2: Querying (Fast Retrieval) ---
 query_template = "Query: {query}"
 query_text = query_template.format(query=query)
-# Now we add EOS to mark the end of the full sequence
 query_tokens = tokenizer.encode(query_text, add_eos=True)
 
-print(f"Processing query: '{query}' against {len(cached_states)} cached docs...")
-
-# We can batch the query processing against multiple document states
-# 1. Prepare a batch of states (Move back to GPU)
-#    Note: We must CLONE/DEEPCOPY because RWKV modifies state in-place!
 batch_states = [[], []]
 for cpu_s in cached_states:
-    batch_states[0].append(cpu_s[0].clone().cuda()) # Tokenshift State
-    batch_states[1].append(cpu_s[1].clone().cuda()) # TimeMix State
+    batch_states[0].append(cpu_s[0].clone().cuda())
+    batch_states[1].append(cpu_s[1].clone().cuda())
 
-# Stack into batch tensors
-# State[0]: [Layers, 2, 1, Hidden] -> Stack dim 2 -> [Layers, 2, Batch, Hidden]
-# State[1]: [Layers, 1, Heads, HeadSize, HeadSize] -> Stack dim 1 -> [Layers, Batch, Heads, ...]
 state_input = [
     torch.stack(batch_states[0], dim=2).squeeze(3), 
     torch.stack(batch_states[1], dim=1).squeeze(2)
 ]
 
-# 2. Prepare query tokens (Broadcast query to batch size)
-batch_size = len(documents)
-batch_query_tokens = [query_tokens] * batch_size
-
-# 3. Fast Forward (Only processing query tokens!)
+batch_query_tokens = [query_tokens] * len(documents)
 _, final_state = emb_model.forward(batch_query_tokens, state_input)
 logits = reranker.forward(final_state[1])
 scores = torch.sigmoid(logits)
-
-print("\nRWKVReRanker Offline Scores:")
-for doc, score in zip(documents, scores):
-    print(f"[{score:.4f}] {doc}")
 ```
 
 ## Summary of Differences
@@ -221,4 +182,15 @@ for doc, score in zip(documents, scores):
 | **Latency** | Extremely Fast | Slow O(L_doc + L_query) | Fast O(L_query) only |
 | **Input** | Query & Doc separate | `Instruct + Doc + Query` | `Query` (on top of cached Doc) |
 | **Storage** | Low (Vector only) | None | High (Stores Hidden States) |
-| **Best For** | Initial Retrieval (Top-k) | Reranking few candidates | Reranking many candidates |
+| **Best For** | Initial Retrieval (Top-k) | Reranking few candidates | Reranking many candidates |
+
+## Citation
+
+```bibtex
+@article{hou2025embeddingrwkv,
+  title={EmbeddingRWKV: State-Centric Retrieval with Reusable States},
+  author={Hou, Howard and others},
+  journal={arXiv preprint arXiv:2601.07861},
+  year={2026}
+}
+```