| """North Micro Vision's visual tower and all deepstack mergers.""" | |
| import json | |
| from pathlib import Path | |
| from accelerate import init_empty_weights | |
| from safetensors.torch import load_file | |
| from transformers.models.cohere_compass.configuration_cohere_compass import CohereCompassVisionConfig | |
| from transformers.models.cohere_compass.modeling_cohere_compass import CohereCompassVisionModel | |
| from transformers.models.cohere_compass.image_processing_cohere_compass import CohereCompassImageProcessor | |
| def load_model(root): | |
| root=Path(root);c=json.loads((root/'config.json').read_text())['vision_config'];config=CohereCompassVisionConfig(**c);config._attn_implementation='eager' | |
| with init_empty_weights():model=CohereCompassVisionModel(config) | |
| state=load_file(str(root/'model.safetensors'),device='cpu');model.load_state_dict({k.removeprefix('model.visual.'):v for k,v in state.items()},strict=True,assign=True) | |
| return model.eval() | |
| def load_processor(root): | |
| config=json.loads((Path(root)/'preprocessor_config.json').read_text());return CohereCompassImageProcessor.from_dict(config) | |