93 lines
2.4 KiB
Markdown
93 lines
2.4 KiB
Markdown
# Usage Guide
|
|
|
|
## Requirements
|
|
|
|
- Python 3.10+
|
|
- transformers >= 4.54.0
|
|
- CUDA-capable GPU with 8GB+ VRAM (for bfloat16 inference)
|
|
|
|
## Installation
|
|
|
|
```bash
|
|
pip install transformers>=4.54.0 torch
|
|
```
|
|
|
|
For better performance with vLLM:
|
|
|
|
```bash
|
|
pip install vllm==0.10.1
|
|
```
|
|
|
|
## Loading the Model
|
|
|
|
```python
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
"OMCHOKSI108/VibeThinker-3B",
|
|
low_cpu_mem_usage=True,
|
|
torch_dtype="bfloat16",
|
|
device_map="auto",
|
|
)
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
"OMCHOKSI108/VibeThinker-3B",
|
|
trust_remote_code=True,
|
|
)
|
|
```
|
|
|
|
## Basic Inference
|
|
|
|
```python
|
|
from transformers import GenerationConfig
|
|
|
|
messages = [{"role": "user", "content": "Solve for x: 3x + 7 = 22"}]
|
|
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
|
|
inputs = tokenizer([text], return_tensors="pt").to(model.device)
|
|
|
|
outputs = model.generate(
|
|
**inputs,
|
|
generation_config=GenerationConfig(
|
|
max_new_tokens=40960,
|
|
do_sample=True,
|
|
temperature=0.6,
|
|
top_p=0.95,
|
|
top_k=None,
|
|
),
|
|
)
|
|
response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
|
|
print(response)
|
|
```
|
|
|
|
## Recommended Parameters
|
|
|
|
| Parameter | Value | Notes |
|
|
|-----------|-------|-------|
|
|
| temperature | 0.6 or 1.0 | Lower for focused answers, higher for diversity |
|
|
| top_p | 0.95 | Nucleus sampling threshold |
|
|
| top_k | None (or -1 in vLLM/SGLang) | Skip top-k filtering |
|
|
| max_new_tokens | 40960-102400 | Longer for complex reasoning tasks |
|
|
|
|
## Using with vLLM
|
|
|
|
```python
|
|
from vllm import LLM, SamplingParams
|
|
|
|
llm = LLM("OMCHOKSI108/VibeThinker-3B", dtype="bfloat16")
|
|
params = SamplingParams(temperature=0.6, top_p=0.95, top_k=-1, max_tokens=40960)
|
|
output = llm.generate("What is the derivative of x^3?", params)
|
|
print(output[0].outputs[0].text)
|
|
```
|
|
|
|
## Hardware Requirements
|
|
|
|
- **Minimum:** 8 GB VRAM (bfloat16, 3B model)
|
|
- **Recommended:** 16 GB+ VRAM (allows larger batch sizes or longer sequences)
|
|
- **CPU inference:** Possible but significantly slower
|
|
|
|
## Important Notes
|
|
|
|
1. This model is optimized for verifiable reasoning tasks (math, code, STEM).
|
|
2. It is not designed for tool-calling, agent-based programming, or function calling.
|
|
3. For open-domain knowledge tasks, larger models may be more suitable.
|
|
4. See the [original model card](./ORIGINAL_README.md) for further guidance from the authors.
|