| """Download and load released Sage Sci Perception projectors.""" |
| import json |
| from pathlib import Path |
|
|
| from safetensors.torch import load_file |
|
|
| from modeling_projector import build |
|
|
|
|
| def download_projector(projector_id, repo_id='internlm/Sage-SciPerception', |
| revision='main', local_dir=None, cache_dir=None): |
| """Resolve a projector through the repository index and download its files.""" |
| from huggingface_hub import hf_hub_download |
|
|
| index_path = Path(hf_hub_download(repo_id, 'config.json', revision=revision, |
| cache_dir=cache_dir)) |
| index = json.loads(index_path.read_text()) |
| component = index['projectors'][projector_id] |
| |
| resolved_revision = index_path.parent.name |
| if local_dir is not None: |
| destination = Path(local_dir) |
| destination.mkdir(parents=True, exist_ok=True) |
| (destination / 'config.json').write_bytes(index_path.read_bytes()) |
| config_path = Path(hf_hub_download( |
| repo_id, component + '/config.json', revision=resolved_revision, |
| local_dir=local_dir, cache_dir=cache_dir, |
| )) |
| config = json.loads(config_path.read_text()) |
| hf_hub_download(repo_id, component + '/' + config['weights'], |
| revision=resolved_revision, local_dir=local_dir, cache_dir=cache_dir) |
| return config_path.parent |
|
|
|
|
| def load_projector(directory, device='cpu'): |
| directory = Path(directory) |
| config = json.loads((directory / 'config.json').read_text()) |
| model = build(config['projector'], config['output_dim']) |
| state = load_file(str(directory / config['weights']), device='cpu') |
| model.load_state_dict(state, strict=True) |
| return model.to(device).requires_grad_(False).eval(), config |
|
|
|
|
| if __name__ == '__main__': |
| import argparse |
|
|
| parser = argparse.ArgumentParser(description='Download a projector selected from the repository index.') |
| parser.add_argument('projector_id', help='backbone/modality/task, as listed in the root config.json') |
| parser.add_argument('--repo-id', default='internlm/Sage-SciPerception') |
| parser.add_argument('--revision', default='main') |
| parser.add_argument('--local-dir', type=Path) |
| args = parser.parse_args() |
| print(download_projector(args.projector_id, repo_id=args.repo_id, |
| revision=args.revision, local_dir=args.local_dir)) |
|
|
