Update README.md

README.md

@@ -1,3 +1,4 @@
+
 ---
 datasets:
 - Aquiles-ai/Athenea-Coding-100k
@@ -8,29 +9,38 @@ tags:
 - merge
 - uncensored
 - Rhea
+- multi-pass
+- reasoning
 language:
 - en
 pipeline_tag: text-generation
 library_name: transformers
 ---
 
-<h1 align="center">Rhea 4BCoding</h1>
+<h1 align="center">Rhea 4B Coding</h1>
 
 ![image](atheneamodel.png)
 
-**Rhea-4B-Coding** is a optimized version of [Aquiles-ai/Athenea-4B-Coding](https://huggingface.co/Aquiles-ai/Athenea-4B-Coding), specialized in **code reasoning, debugging, agentic tools and problem solving**.
-Trained on high-quality programming data with explicit reasoning traces using `<think>` and `</think>` tags, the model is designed to perform detailed step-by-step reasoning for software development, algorithm design, and code comprehension tasks.
+**Rhea-4B-Coding** is an optimized version of [Aquiles-ai/Athenea-4B-Coding](https://huggingface.co/Aquiles-ai/Athenea-4B-Coding), specialized in **code reasoning, debugging, agentic tools and multi-pass problem solving**.
+
+Trained on high-quality programming data with explicit reasoning traces using ` thinking` and ` 思考结束` tags, the model is designed to perform detailed **3-pass reasoning** for software development, algorithm design, and code comprehension tasks:
+
+1. **Pass 1**: First implementation
+2. **Pass 2**: Self-review for bugs, edge cases, security, performance
+3. **Pass 3**: Final optimized version with identical functionality
 
 > ⚠️ **Important Note:** This model uses an *uncensored* base version, providing full expressive freedom and unrestricted output generation. Users are fully responsible for any use or content produced by the model. It is intended exclusively for research and experimentation purposes.
 
 ## 🎯 Model Description
 
-Rhea-4B-Coding extends Athenea-4B-Coding’s structured reasoning capabilities into programming-related domains, showing strong performance on logical problem-solving, code completion, and debugging scenarios.
+Rhea-4B-Coding extends Athenea-4B-Coding's structured reasoning capabilities into programming-related domains with **multi-pass processing**, showing strong performance on logical problem-solving, code completion, debugging scenarios, and iterative code refinement.
 
 Key features:
 
+* **Multi-Pass Processing**: 3-step reasoning with `<think>`, `<review>`, and `<final>` tags
 * **Agentic Tools** for AI Agents
-* **Step-by-step code reasoning** within `<think>` blocks
+* **Step-by-step code reasoning** within ` thinking` blocks
+* **Self-review capabilities** for bug detection and optimization
 * **Specialization in algorithmic and debugging tasks**
 * **Uncensored output generation** for full reasoning visibility
 * **Improved logical consistency** through focused fine-tuning
@@ -38,7 +48,20 @@ Key features:
 
 The model was fine-tuned using the dataset [Aquiles-ai/Athenea-Coding-100k](https://huggingface.co/datasets/Aquiles-ai/Athenea-Coding-100k), which includes diverse programming challenges, structured reasoning chains, and natural language explanations across multiple programming languages.
 
+## 🔄 Multi-Pass Architecture
+
+The model uses special tokens for structured reasoning:
 
+| Token | Purpose |
+|-------|---------|
+| `<think>` | Start of self-review phase (Pass 2) |
+| `</think>` | End of self-review phase |
+| `<review>` | Start of review results documentation |
+| `</review>` | End of review results |
+| `<final>` | Start of final optimized version (Pass 3) |
+| `</final>` | End of final version |
+
+This structure ensures **identical functionality** across all passes while improving **code structure, comments, and robustness**.
 
 ## 💻 Usage
 
@@ -83,13 +106,101 @@ inputs = tokenizer.apply_chat_template(
 with torch.no_grad():
     output = model.generate(
         **inputs,
-        max_new_tokens=8092,
-        pad_token_id=tokenizer.eos_token_id,
+        max_new_tokens=16384,  # Increased for multi-pass output
+        pad_token_id=tokenizer.pad_token_id,
         eos_token_id=tokenizer.eos_token_id,
     )
 
 # Decode and print the output
-print(tokenizer.decode(output[0], skip_special_tokens=True))
+print(tokenizer.decode(output[0], skip_special_tokens=False))
+```
+
+### Multi-Pass Inference (Recommended)
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+model = AutoModelForCausalLM.from_pretrained("Roskosmos19/Rhea-4B-Coding",
+        dtype=torch.bfloat16,
+        trust_remote_code=True,
+        device_map="auto")
+
+tokenizer = AutoTokenizer.from_pretrained("Roskosmos19/Rhea-4B-Coding", trust_remote_code=True)
+
+def generate_multi_pass(prompt, max_tokens_per_pass=4096):
+    """
+    Generate code with 3-pass reasoning:
+    Pass 1: First implementation
+    Pass 2: Self-review
+    Pass 3: Final optimized version
+    """
+    
+    # Pass 1: First implementation
+    messages = [{"role": "user", "content": prompt}]
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt",
+    ).to('cuda')
+    
+    with torch.no_grad():
+        output1 = model.generate(
+            **inputs,
+            max_new_tokens=max_tokens_per_pass,
+            temperature=0.4,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+    
+    pass1 = tokenizer.decode(output1[0], skip_special_tokens=False)
+    
+    # Pass 2: Self-review
+    review_prompt = pass1 + "\n<<think>\n### PASS 2 - Self-Review:\n"
+    inputs2 = tokenizer(review_prompt, return_tensors="pt").to('cuda')
+    
+    with torch.no_grad():
+        output2 = model.generate(
+            **inputs2,
+            max_new_tokens=2048,
+            temperature=0.3,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+    
+    review = tokenizer.decode(output2[0], skip_special_tokens=False)
+    
+    # Pass 3: Final version
+    final_prompt = review + "\n<<final>\n### PASS 3 - Final Version:\n"
+    inputs3 = tokenizer(final_prompt, return_tensors="pt").to('cuda')
+    
+    with torch.no_grad():
+        output3 = model.generate(
+            **inputs3,
+            max_new_tokens=max_tokens_per_pass,
+            temperature=0.2,
+            pad_token_id=tokenizer.pad_token_id,
+            eos_token_id=tokenizer.eos_token_id,
+        )
+    
+    final = tokenizer.decode(output3[0], skip_special_tokens=False)
+    
+    return {
+        "pass1": pass1,
+        "review": review,
+        "pass3": final
+    }
+
+# Example usage
+result = generate_multi_pass("Write a Python function for binary search")
+print("=== PASS 1 ===")
+print(result["pass1"])
+print("\n=== REVIEW ===")
+print(result["review"])
+print("\n=== FINAL ===")
+print(result["pass3"])
 ```
 
 ### Streaming Inference
@@ -99,13 +210,13 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStream
 import torch
 from threading import Thread
 
-model = AutoModelForCausalLM.from_pretrained("Aquiles-ai/Athenea-4B-Coding",
+model = AutoModelForCausalLM.from_pretrained("Roskosmos19/Rhea-4B-Coding",
         dtype=torch.bfloat16,
         trust_remote_code=True,
         device_map="auto",
         attn_implementation="flash_attention_2")
 
-tokenizer = AutoTokenizer.from_pretrained("Aquiles-ai/Athenea-4B-Coding", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("Roskosmos19/Rhea-4B-Coding", trust_remote_code=True)
 
 messages = [
     {"role": "user", "content": "Hey, write a Python function that implements the binary search algorithm recursively."}
@@ -120,13 +231,13 @@ inputs = tokenizer.apply_chat_template(
 ).to('cuda')
 
 # Create the streamer
-streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=False)
 
 # Build kwargs for generate
 generate_kwargs = dict(
     **inputs,
-    max_new_tokens=8092,
-    pad_token_id=tokenizer.eos_token_id,
+    max_new_tokens=16384,  # Increased for multi-pass output
+    pad_token_id=tokenizer.pad_token_id,
     eos_token_id=tokenizer.eos_token_id,
     streamer=streamer,
 )
@@ -151,7 +262,7 @@ vllm serve Roskosmos19/Rhea-4B-Coding \
   --host 0.0.0.0 \
   --port 8000 \
   --api-key dummyapikey \
-  --max-model-len=16384 \
+  --max-model-len=262144 \
   --async-scheduling \
   --gpu-memory-utilization=0.90
 ```
@@ -167,7 +278,7 @@ stream = client.chat.completions.create(
         "role": "user",
         "content": "Hey, write a Python function that determines if a string is a palindrome, ignoring case, spaces, and punctuation."
     }],
-    max_tokens=8092,
+    max_tokens=16384,  # Increased for multi-pass output
     stream=True
 )
 for chunk in stream:
@@ -177,8 +288,32 @@ for chunk in stream:
 
 **vLLM Benefits:** 20-30x faster inference, OpenAI-compatible API, continuous batching, async scheduling.
 
+## 📝 Model Configuration
+
+| Parameter | Value | Description |
+|-----------|-------|-------------|
+| `temperature` | 0.4 | Balanced creativity and consistency |
+| `max_new_tokens` | 32768 | Full multi-pass output capacity |
+| `repetition_penalty` | 1.0 | No penalty for intentional code repetition |
+| `no_repeat_ngram_size` | 0 | Allows code structure repetition |
+| `use_cache` | true | Faster inference for long outputs |
+
+## ⚙️ Files Modified for Multi-Pass
+
+| File | Changes |
+|------|---------|
+| `generation_config.json` | Extended tokens, optimized for multi-pass |
+| `config.json` | Enabled caching, full context window |
+| `tokenizer_config.json` | Added `<think>`, `<review>`, `<final>` tokens |
+| `special_tokens_map.json` | Registered new special tokens |
+| `chat_template.jinja` | 3-pass prompt structure |
+
+## 🤝 Credits
 
+- Base model: [Aquiles-ai/Athenea-4B-Coding](https://huggingface.co/Aquiles-ai/Athenea-4B-Coding)
+- Dataset: [Aquiles-ai/Athenea-Coding-100k](https://huggingface.co/datasets/Aquiles-ai/Athenea-Coding-100k)
+- Architecture: Qwen3 4B
 
 <p align="center">
   Roskosmos19
-</p>
+</p>