StyTJU/robobrain2.0-3b-mor-router-adapter

RoboBrain-3B MoR 稀疏激活推理使用说明

本仓库（StyTJU/robobrain2.0-3b-mor-router-adapter）为 RoboBrain2.0-3B 的 MoR（Mixture-of-Recursions）稀疏激活路由器+lm_head 增量权重，需配合官方基座模型使用。

1. 环境准备

按照BAAI/RoboBrain2.0-3B的环境配置即可

pip install torch transformers safetensors

2. 加载基座模型与 MoR 适配器权重

from robobrain3b_mor_model import load_robobrain3b_with_mor
from transformers import AutoTokenizer, AutoProcessor
import torch, os, json
from safetensors.torch import load_file

# 1. 加载基座模型（官方原版）
model = load_robobrain3b_with_mor(
    model_path="BAAI/RoboBrain2.0-3B",
    device_map="auto",
    torch_dtype=torch.bfloat16
)

# 2. 加载MoR适配器权重（本仓库内容）
adapter_dir = "./outputs/mor-kv-sharerobot"  # 如 './outputs/mor-kv-sharerobot' 或 HF下载目录
index_file = os.path.join(adapter_dir, 'model.safetensors.index.json')
index = json.load(open(index_file, 'r', encoding='utf-8'))
weight_map = index["weight_map"]
from collections import defaultdict
grouped = defaultdict(list)
for k, v in weight_map.items():
    grouped[v].append(k)
for shard, keys in grouped.items():
    tens = load_file(os.path.join(adapter_dir, shard))
    sub = {k: tens[k] for k in keys if k in tens}
    model.load_state_dict(sub, strict=False)
model.to(model.device)
print("✅ MoR适配器权重加载完成")

3. 加载分词器与处理器

tokenizer = AutoTokenizer.from_pretrained("BAAI/RoboBrain2.0-3B", trust_remote_code=True)
processor = AutoProcessor.from_pretrained("BAAI/RoboBrain2.0-3B", trust_remote_code=True)

4. 多模态推理示例（图片+文本）

messages = [
    {
        "role": "user",
        "content": [
            {
                "type": "image",
                "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg",
            },
            {"type": "text", "text": "请描述这张图片。"},
        ],
    }
]

# 构建对话模板
text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
image_inputs = ["https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"]
inputs = processor(
    text=[text],
    images=image_inputs,
    padding=True,
    return_tensors="pt",
)
inputs = {k: v.to(model.device) for k, v in inputs.items()}

# 推理，可调整超参数
generated_ids = model.generate(
    **inputs,
    max_new_tokens=128,
    temperature=0.2,   # 控制多样性
    top_p=0.95,        # nucleus采样
    do_sample=True,    # 启用采样
    pad_token_id=tokenizer.eos_token_id,
)
in_len = inputs['input_ids'].shape[1]
output_text = tokenizer.decode(generated_ids[0][in_len:], skip_special_tokens=True)
print(output_text)