vLLM

bash
Copy code
export MODEL_DIR=/path/to/TranR-14B
export PYTHONPATH="$MODEL_DIR/runtime_patches/qwen_rnn:$PYTHONPATH"
export QWEN_RNN_ENABLED=1
export QWEN_RNN_MODEL_DIR="$MODEL_DIR"

python -m vllm.entrypoints.openai.api_server \
  --model "$MODEL_DIR" \
  --served-model-name TranR-14B \
  --tensor-parallel-size 4 \
  --max-model-len 40960 \
  --dtype auto \
  --trust-remote-code \
  --enable-reasoning \
  --reasoning-parser deepseek_r1

Transformers

python
Copy code
import os
import sys
import torch

model_dir = "/path/to/TranR-14B"
os.environ["QWEN_RNN_ENABLED"] = "1"
os.environ["QWEN_RNN_MODEL_DIR"] = model_dir
sys.path.insert(0, f"{model_dir}/runtime_patches/qwen_rnn")
import sitecustomize

from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_dir,
    torch_dtype="auto",
    device_map="auto",
    trust_remote_code=True,
)

messages = [{"role": "user", "content": "Solve: 1+1="}]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
    enable_thinking=True,
)
inputs = tokenizer(text, return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=32768)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))