Update README.md

README.md

@@ -1,63 +1,127 @@
-
 ---
 license: apache-2.0
 ---
 
-## Introduction
-
-**InfiR2-1.5B-base-FP8** is derived from the **Qwen2.5-1.5B-base** model via continued pre-training using the **FP8** recipe.
-
-## Model Download
-
-```bash
-# Create a directory for models
-mkdir -p ./models
-# Download InfiR2-1.5B-base-FP8 model
-huggingface-cli download --resume-download InfiX-ai/InfiR2-1.5B-base-FP8 --local-dir ./models/InfiR2-1.5B-base-FP8
-````
-
-
-## Quick Start
+# InfiR2-1.5B-base-FP8
+
+<p align="center">
+  <a href="https://arxiv.org/abs/2509.22536">📄 Paper</a> &nbsp; | &nbsp;
+  <a href="https://github.com/InfiXAI/InfiR2"> 🐙 Github</a> &nbsp; |
+  <a href="https://infix-ai.com/research/infir2/">🌐 Project Website</a> &nbsp;
+</p>
+
+We performed continual pre-training (CPT) on the **Qwen2.5-1.5B-base** model for an additional 160 billion tokens using the FP8 format. In this process, both the forward and backward passes employed the E4M3 format, and quantization scaling factors were represented in UE8M0. The training data mixture was composed of:
+- 140B tokens from public sources, including FineWeb, Nemotron Datasets and stack-edu and issues-kaggle-notebooks.
+- A subsequent 20B tokens mixed with data from AM-DeepSeek-R1 and AM-Qwen3.
+
+The resulting model is the **InfiR2-1.5B-base-FP8**.
+
+
+**Training Recipe**:
+<p align="center">
+    <img src="fp8_recipe.png" width="100%"/>
+<p>
+
+- Stable and Reproducible Performance
+- Efficient and Low memory Training 
+
+
+
+## 🚀 InfiR2 Model Series
+
+The InfiR2 framework offers multiple variants model with different size and training strategy:
+
+- **1.5B**
+- [InfiR2-1.5B-base-FP8](https://huggingface.co/InfiX-ai/InfiR2-1.5B-base-FP8): *Continue pretrain on Qwen2.5-1.5B-base*
+- [InfiR2-1.5B-Instruct-FP8](https://huggingface.co/InfiX-ai/InfiR2-1.5B-Instruct-FP8): *Supervised fine-tuning on InfiR2-1.5B-base-FP8 with [InfiAlign dataset](https://huggingface.co/papers/2508.05496)*
+- **7B**
+- [InfiR2-7B-base-FP8](https://huggingface.co/InfiX-ai/InfiR2-7B-base-FP8): *Continue pretrain on Qwen2.5-7B-base*
+- [InfiR2-7B-Instruct-FP8](https://huggingface.co/InfiX-ai/InfiR2-7B-Instruct-FP8): *Supervised fine-tuning on InfiR2-7B-base-FP8 with [InfiAlign dataset](https://huggingface.co/papers/2508.05496)*
+- [InfiR2-R1-7B-FP8](https://huggingface.co/InfiX-ai/InfiR2-R1-7B-FP8): *Reinforcement learning on InfiR2-7B-Instruct-FP8 with dapo dataset*
+
+## 📊 Model Performance
+The **InfiR2-1.5B-Instruct-FP8** model is the result of further fine-tuning applied to the **InfiR2-1.5B-base-FP8**. For further details, refer to [InfiR2-1.5B-Instruct-FP8](https://huggingface.co/InfiX-ai/InfiR2-1.5B-Instruct-FP8). Below is the performance comparison of InfiR2-1.5B-Instruct-FP8 on reasoning benchmarks. Note: 'w. InfiAlign' denotes Supervised Fine-Tuning (SFT) using the InfiAlign dataset. 
+
+<div align="center">
+
+<table>
+  <thead>
+    <tr>
+      <th align="left">Model</th>
+      <th align="center">AIME 25</th>
+      <th align="center">AIME 24</th>
+      <th align="center">GPQA</th>
+      <th align="center">LiveCodeBench v5</th>
+    </tr>
+  </thead>
+  <tbody>
+    <tr>
+      <td align="left"><strong>Deepseek-Distill-Qwen-1.5B</strong></td>
+      <td align="center">21.35</td>
+      <td align="center">26.87</td>
+      <td align="center">32.26</td>
+      <td align="center">18.50</td>
+    </tr>
+    <tr>
+      <td align="left"><strong>Qwen2.5-1.5B-base (w. InfiAlign)</strong></td>
+      <td align="center">14.58</td>
+      <td align="center">10.52</td>
+      <td align="center">28.98</td>
+      <td align="center">12.99</td>
+    </tr>
+    <tr>
+      <td align="left"><strong>InfiR2-1.5B-Instruct-FP8</strong></td>
+      <td align="center">18.45</td>
+      <td align="center">17.39</td>
+      <td align="center">29.48</td>
+      <td align="center">17.10</td>
+    </tr>
+  </tbody>
+</table>
+
+</div>
+
+
+## 🎭 Quick Start
 
 ```python
+from vllm import LLM, SamplingParams
 import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
+import os
 
 MODEL_NAME = "InfiX-ai/InfiR2-1.5B-base-FP8"
 
 prompt_text = "Briefly explain what a black hole is, and provide two interesting facts."
 
-MAX_NEW_TOKENS = 256  
-TEMPERATURE = 0.8     
-DO_SAMPLE = True      
+MAX_NEW_TOKENS = 256
+TEMPERATURE = 0.8
+DO_SAMPLE = True
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+llm = LLM(
+    model=MODEL_NAME, 
+    dtype="auto", 
+)
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    torch_dtype=torch.bfloat16 if device == "cuda" else None 
-).to(device)
+sampling_params = SamplingParams(
+    n=1,
+    temperature=TEMPERATURE,
+    max_tokens=MAX_NEW_TOKENS,
+)
 
+tokenizer = llm.get_tokenizer()
 messages = [
     {"role": "user", "content": prompt_text}
 ]
-input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt", tokenize=True).to(device) # Added tokenize=True
+prompt_formatted = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
 
-with torch.no_grad():
-    output_ids = model.generate(
-        input_ids,
-        max_new_tokens=MAX_NEW_TOKENS,
-        temperature=TEMPERATURE,
-        do_sample=DO_SAMPLE,
-        pad_token_id=tokenizer.eos_token_id
-    )
+outputs = llm.generate(
+    prompt_formatted, 
+    sampling_params
+)
 
-generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+generated_text = outputs[0].outputs[0].text
 
-# For a base model, the output is often a direct continuation of the prompt.
-response_start_index = generated_text.rfind(prompt_text) + len(prompt_text)
-llm_response = generated_text[response_start_index:].strip()
+llm_response = generated_text.strip()
 
 print("\n" + "="*70)
 print(f"Prompt: \n{prompt_text}")
@@ -66,11 +130,37 @@ print(f"(LLM Response): \n{llm_response}")
 print("="*70)
 ```
 
-## Acknowledgements
+## 📚 Model Download
+
+```bash
+# Create a directory for models
+mkdir -p ./models
+# Download InfiR2-1.5B-base-FP8 model
+huggingface-cli download --resume-download InfiX-ai/InfiR2-1.5B-base-FP8 --local-dir ./models/InfiR2-1.5B-base-FP8
+```
+## 🎯 Intended Uses
+
+### ✅ Direct Use
+
+This model is intended for research and commercial use. Example use cases include:
+
+- Instruction following
+- Mathematical reasoning
+- Code generation
+- General reasoning
+
+### ❌ Out-of-Scope Use
+
+The model should **not** be used for:
+
+- Generating harmful, offensive, or inappropriate content
+- Creating misleading information
+
+## 🙏 Acknowledgements
 
   * We would like to express our gratitude for the following open-source projects: [Slime](https://github.com/THUDM/slime), [Megatron](https://github.com/NVIDIA/Megatron-LM), [TransformerEngine](https://github.com/NVIDIA/TransformerEngine) and [Qwen2.5](https://github.com/QwenLM/Qwen2.5-Math).
 
-## Citation
+## 📌 Citation
 
 If you find our work useful, please cite: