File size: 1,851 Bytes
a86e230 7c14598 7adffc3 7c14598 a86e230 7c14598 a86e230 7c14598 a86e230 7c14598 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | import argparse
import json
from pathlib import Path
from huggingface_hub import snapshot_download
from peft import PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer
ROOT_FILES = [
'adapter_config.json',
'adapter_model.safetensors',
'tokenizer.json',
'tokenizer_config.json',
'chat_template.jinja',
]
def load_adapter(subfolder=None, revision=None):
"""Load the root MuSiQue 1.7B SLS adapter, or a checkpoint subfolder."""
repo_id = 'SJTUFreshman/collaborativelearning'
subfolder = (subfolder or '').replace('\\', '/').strip('/')
if subfolder == '.':
subfolder = ''
tokenizer_folder = str(Path(subfolder).parent).replace('\\', '/') if subfolder.endswith('/policy_adapter') else subfolder
patterns = (
[subfolder+'/*', tokenizer_folder+'/tokenizer*', tokenizer_folder+'/chat_template.jinja', tokenizer_folder+'/special_tokens_map.json']
if subfolder else ROOT_FILES
)
snapshot = snapshot_download(repo_id, revision=revision, allow_patterns=patterns)
adapter_dir = Path(snapshot)/subfolder
config = json.loads((adapter_dir/'adapter_config.json').read_text())
tokenizer = AutoTokenizer.from_pretrained(Path(snapshot)/tokenizer_folder)
base = AutoModelForCausalLM.from_pretrained(config['base_model_name_or_path'], torch_dtype='auto', device_map='auto')
model = PeftModel.from_pretrained(base, adapter_dir)
return model, tokenizer
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('subfolder', nargs='?', help='Checkpoint directory; default: root MuSiQue 1.7B SLS adapter')
parser.add_argument('--revision')
args = parser.parse_args()
model, tokenizer = load_adapter(args.subfolder, args.revision)
print('Loaded', args.subfolder or 'root (MuSiQue/SLS/Qwen3-1.7B/final)')
|