Abdoulaye coumbassa
openhands
feat(vision): trainable stage-1 vision projector with frozen encoder and LM
ba3bbc7 | """Stage-1 vision projector training for CORTEX AI. | |
| Trains *only* ``VisionProjector`` on top of a frozen real vision encoder and a frozen small | |
| causal LM:: | |
| IMAGE -> real vision encoder -> VisionProjector -> PrefixFusion -> frozen causal LM -> logits | |
| Nothing here touches the DeepSeek-V4-Pro shards: the ``model-*-of-*.safetensors`` files are only | |
| read to hash their LFS pointers, and the checkpoint written by this script is a separate file. | |
| Examples:: | |
| python examples/train_vision_projector.py --config configs/vision_projector.json | |
| python examples/train_vision_projector.py --max-steps 20 --batch-size 2 | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import sys | |
| from dataclasses import asdict, fields | |
| from pathlib import Path | |
| import torch | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) | |
| from cortex_ai.vision.train import ( # noqa: E402 | |
| ProjectorTrainingConfig, | |
| build_fusion_pipeline, | |
| count_parameters, | |
| iterate_caption_batches, | |
| run_training, | |
| save_projector, | |
| ) | |
| def load_dataset(config: ProjectorTrainingConfig) -> tuple[list, list[str]]: | |
| """Load real images and captions from the configured Hub dataset.""" | |
| from datasets import load_dataset as hf_load_dataset | |
| ds = hf_load_dataset(config.dataset, split=config.dataset_split, streaming=True) | |
| images, captions = [], [] | |
| limit = config.limit or 0 | |
| for example in ds: | |
| image = example.get("jpg") or example.get("image") | |
| text = example.get("txt") or example.get("caption") or example.get("text") | |
| if image is None or not text: | |
| continue | |
| if isinstance(text, list): | |
| text = text[0] | |
| images.append(image.convert("RGB")) | |
| captions.append(str(text)) | |
| if limit and len(images) >= limit: | |
| break | |
| if not images: | |
| raise SystemExit(f"no usable image/caption pairs found in {config.dataset}") | |
| return images, captions | |
| def parse_args(argv: list[str] | None = None) -> ProjectorTrainingConfig: | |
| p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) | |
| p.add_argument("--config", help="JSON file of ProjectorTrainingConfig fields") | |
| defaults = ProjectorTrainingConfig() | |
| for f in fields(defaults): | |
| value = getattr(defaults, f.name) | |
| if isinstance(value, bool): | |
| p.add_argument(f"--{f.name.replace('_', '-')}", type=lambda s: s.lower() == "true", default=None) | |
| else: | |
| p.add_argument( | |
| f"--{f.name.replace('_', '-')}", | |
| type=type(value) if value is not None else str, | |
| default=None, | |
| ) | |
| args = p.parse_args(argv) | |
| values = {} | |
| if args.config: | |
| values.update(json.loads(Path(args.config).read_text())) | |
| for f in fields(defaults): | |
| v = getattr(args, f.name.replace("-", "_")) | |
| if v is not None: | |
| values[f.name] = v | |
| return ProjectorTrainingConfig(**values) | |
| def main(argv: list[str] | None = None) -> int: | |
| config = parse_args(argv) | |
| torch.manual_seed(config.seed) | |
| print("=" * 78) | |
| print("CORTEX AI - stage-1 vision projector training (experimental)") | |
| print("=" * 78) | |
| print(f"config : {json.dumps(asdict(config), indent=2)}") | |
| from transformers import AutoImageProcessor, AutoModelForCausalLM, AutoTokenizer, SiglipVisionModel | |
| dtype = config.resolved_dtype() | |
| device = config.resolved_device() | |
| print(f"\n[1] loading real vision encoder : {config.vision_model}") | |
| image_processor = AutoImageProcessor.from_pretrained(config.vision_model) | |
| vision_encoder = SiglipVisionModel.from_pretrained(config.vision_model).to(dtype=dtype, device=device) | |
| print(f"[2] loading real frozen LM : {config.language_model}") | |
| core = AutoModelForCausalLM.from_pretrained(config.language_model, dtype=dtype).to(device) | |
| tokenizer = AutoTokenizer.from_pretrained(config.language_model) | |
| if tokenizer.pad_token is None: | |
| tokenizer.pad_token = tokenizer.eos_token | |
| print("[3] assembling fusion pipeline (encoder + LM frozen, projector trainable)") | |
| fusion, vision_hidden, lm_hidden = build_fusion_pipeline( | |
| vision_encoder, core, config, image_processor=image_processor | |
| ) | |
| projector_params = count_parameters(fusion.projector) | |
| print(f" vision_hidden_size = {vision_hidden}") | |
| print(f" lm_hidden_size = {lm_hidden}") | |
| print(f" projector_type = {config.projector_type}") | |
| print(f" projector_parameters = {projector_params:,}") | |
| print(f" dtype / device = {config.dtype} / {config.device}") | |
| print(f"\n[4] loading real dataset : {config.dataset} (split={config.dataset_split})") | |
| images, captions = load_dataset(config) | |
| print(f" {len(images)} image/caption pairs") | |
| print(f" sample caption: {captions[0][:70]!r}") | |
| batches = iterate_caption_batches(images, captions, tokenizer, image_processor, config) | |
| def on_step(step: int, loss: float) -> None: | |
| if step == 1 or step % config.log_every == 0: | |
| print(f" step {step:4d} loss={loss:.4f}") | |
| print(f"\n[5] training projector for up to {config.max_steps} steps (projector only)") | |
| report = run_training( | |
| fusion, batches, config, projector_params, vision_hidden, lm_hidden, on_step=on_step | |
| ) | |
| path = save_projector(fusion.projector, config, report) | |
| report.checkpoint_path = str(path) | |
| print("\n" + "=" * 78) | |
| print("RESULT") | |
| print("=" * 78) | |
| print(report.to_json()) | |
| print(f"\ncheckpoint written: {path}") | |
| print("\nPROVEN by this run:") | |
| print(" - real image -> real encoder -> projector -> PrefixFusion -> frozen LM executed") | |
| print(" - projector received gradients and its parameters changed") | |
| print(" - vision encoder and causal LM parameters are unchanged") | |
| print("\nNOT PROVEN:") | |
| print(" - that this projector works with the real DeepSeek-V4-Pro core") | |
| print(" - native multimodal understanding of the 1.65T checkpoint") | |
| return 0 | |
| if __name__ == "__main__": | |
| raise SystemExit(main()) |