Views
No views yet
Axion-Flash-Reasoning-2B is a fine-tuned version of NVIDIA's state-of-the-art Nemotron-Research-Reasoning-Qwen-1.5B model. This version is specifically adapted to be more instruction-friendly and computationally efficient, making it ideal for integration into applications requiring powerful reasoning capabilities without the overhead of larger models.cc-by-nc-4.0)transformers library.pipelinetext-generation pipeline.1from transformers import pipeline
2import torch
3
4# For optimal performance, use a GPU
5pipe = pipeline(
6 "text-generation",
7 model="AdvRahul/Axion-Flash-Reasoning-2B",
8 torch_dtype=torch.bfloat16,
9 device_map="auto"
10)
11
12# Qwen models use a specific chat template. The pipeline handles this automatically.
13messages = [
14 {"role": "system", "content": "You are a helpful assistant that excels at logical reasoning."},
15 {"role": "user", "content": "I have 3 apples and I buy 5 more. I then give 2 apples to my friend. How many apples do I have left?"}
16]
17
18prompt = pipe.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
19outputs = pipe(prompt, max_new_tokens=256, do_sample=True, temperature=0.7, top_k=50, top_p=0.95)
20
21print(outputs[0]["generated_text"])
22
23Optimized Inference (4-bit Quantization)
24To achieve "flash" speed and reduce memory usage, you can load the model in 4-bit using bitsandbytes.
25
26Bash
27
28pip install transformers torch accelerate bitsandbytes
29Python
30
31import torch
32from transformers import AutoModelForCausalLM, AutoTokenizer
33
34model_id = "AdvRahul/Axion-Flash-Reasoning-2B"
35tokenizer = AutoTokenizer.from_pretrained(model_id)
36model = AutoModelForCausalLM.from_pretrained(
37 model_id,
38 torch_dtype=torch.bfloat16,
39 device_map="auto",
40 # This enables 4-bit quantization
41 load_in_4bit=True
42)
43
44messages = [
45 {"role": "system", "content": "You are an expert code assistant."},
46 {"role": "user", "content": "Write a Python function to calculate the factorial of a number using recursion."}
47]
48prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
49inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
50
51outputs = model.generate(**inputs, max_new_tokens=150)
52print(tokenizer.decode(outputs[0], skip_special_tokens=True))
53📝 Model Description
54Fine-Tuning Philosophy
55While the base Nemotron-Research-Reasoning model demonstrates world-class capabilities in formal reasoning (math, code, logic), Axion-Flash has been further instruction-tuned to make these powerful abilities more accessible and practical for real-world applications. The goal is to bridge the gap between a pure research model and a deployable, instruction-following assistant that developers can easily integrate into their products.
56
57This fine-tuning enhances the model's ability to understand and follow user instructions in a conversational format, unlocking its reasoning power for a broader range of tasks.
58
59Key Capabilities
60Complex Reasoning: Inherits the base model's strength in solving logic puzzles, scientific questions, and multi-step problems.
61
62Code Generation: Proficient in generating code for various programming challenges and tasks.
63
64Mathematical Prowess: Excels at solving mathematical problems, from basic arithmetic to more complex Olympiad-level questions.
65
66Enhanced Instruction Following: Fine-tuned to better adhere to user instructions and constraints in a chat-like setting.
67
68ℹ️ Base Model Information (Nemotron-Research-Reasoning-Qwen-1.5B)
69<details>
70<summary>Click to expand details on the powerful base model</summary>
71
72Nemotron-Research-Reasoning-Qwen-1.5B is a leading open-weight model for complex reasoning, trained by NVIDIA using the ProRL (Prolonged Reinforcement Learning) algorithm. This advanced training method enables the model to explore reasoning strategies more deeply, leading to significant performance gains.
73
74The base model was trained on a diverse set of datasets, including:
75
76DeepScaleR-Preview-Dataset
77
78Eurus-2-RL-Data
79
80Reasoning-gym
81
82IFEval
83
84SCP-116K
85
86It sets a new state-of-the-art standard for models in its size class, outperforming competitors by a large margin on benchmarks for math, coding, logic puzzles, and STEM reasoning. For detailed performance metrics, please refer to the original model card.
87
88</details>
89
90⚖️ License and Terms of Use
91This model is released under the cc-by-nc-4.0 license, inheriting the license of its base model.
92
93This means it is available for research and non-commercial use only. Please review the license terms before using this model in your projects.