Views
No views yet
transformers is transformers=4.50.0.1import os
2import json
3import torch
4import argparse
5from transformers import (
6 AutoConfig,
7 AutoModelForCausalLM,
8 AutoTokenizer,
9)
10
11# register the FineRMoE
12def hf_register(hf_ckpt_path):
13 has_py_file = False
14 for filename in os.listdir(hf_ckpt_path):
15 if filename.endswith(".py"):
16 has_py_file = True
17 modeling_filename = filename.split('.')[0]
18 if has_py_file:
19 print("There exists a modeling file.")
20 config_file = os.path.join(hf_ckpt_path, 'config.json')
21 with open(config_file, "r", encoding="utf-8") as file:
22 config = json.load(file)
23 model_type = config['model_type']
24
25 import sys, inspect, importlib
26 sys.path.append(hf_ckpt_path)
27 module = importlib.import_module(modeling_filename)
28
29 for name, obj in inspect.getmembers(module, inspect.isclass):
30 if not obj.__module__ == module.__name__:
31 continue
32 print(name, obj)
33 if name.endswith("CausalLM"):
34 print(f"Found CausalLM with name: {name}")
35 model_module = obj
36 elif name.endswith("Config"):
37 print(f"Found Config with name: {name}")
38 config_module = obj
39
40 AutoConfig.register(model_type, config_module)
41 AutoModelForCausalLM.register(config_module, model_module)
42
43@torch.inference_mode()
44def main(args):
45 # load the tokenizer and the model
46 hf_register(args.model)
47 tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
48 model = AutoModelForCausalLM.from_pretrained(
49 args.model,
50 torch_dtype=torch.float16,
51 attn_implementation=None,
52 device_map="auto",
53 trust_remote_code=True,
54 )
55 device = model.device if hasattr(model, "device") else next(model.parameters()).device
56
57 # prepare the model input
58 prompt = "Please introduce the history of artificial intelligence."
59 messages = [
60 {"role": "user", "content": prompt}
61 ]
62 text = tokenizer.apply_chat_template(
63 messages,
64 tokenize=False,
65 add_generation_prompt=True
66 )
67 inputs = tokenizer([text], return_tensors="pt").to(device)
68
69 # conduct text completion
70 outputs = model.generate(
71 **inputs,
72 max_new_tokens=32000,
73 do_sample=False,
74 use_cache=True,
75 pad_token_id=tokenizer.eos_token_id,
76 )
77
78 output_ids = outputs[0][len(inputs.input_ids[0]):].tolist()
79 content = tokenizer.decode(output_ids, skip_special_tokens=True).strip("\n")
80 print("content:", content)
81
82if __name__ == "__main__":
83 parser = argparse.ArgumentParser()
84 parser.add_argument("--model", default="FineRMoE-26.65B-A7.94B")
85 args = parser.parse_args()
86 main(args)@misc{liao2026finermoedimensionexpansionfinergrained,
title={FineRMoE: Dimension Expansion for Finer-Grained Expert with Its Upcycling Approach},
author={Ning Liao and Xiaoxing Wang and Xiaohan Qin and Junchi Yan},
year={2026},
eprint={2603.13364},
archivePrefix={arXiv},
primaryClass={cs.CV},
url={https://arxiv.org/abs/2603.13364},
}