| 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)') |
|
|
