| """Load KORMo-Web-10B with transformers 5. |
| |
| KORMo-Web-10B is a LLaVA-OneVision model (SigLIP vision tower + projector) with the KORMo-10B text model. |
| transformers does not ship the KORMo text model, so two small steps are needed: |
| |
| import kormo_load |
| kormo_load.register(MODEL_DIR) # 1. register the KORMo text model from this repo |
| model = AutoModelForImageTextToText.from_pretrained(MODEL_DIR, dtype="bfloat16") |
| kormo_load.patch_vision(model) # 2. use the post-layernorm SigLIP features, as in training |
| """ |
| from transformers import AutoConfig, AutoModel, AutoModelForCausalLM |
| from transformers.dynamic_module_utils import get_class_from_dynamic_module |
|
|
|
|
| def register(model_dir): |
| """Register `kormo` (from `_configuration_kormo.py` / `_modeling_kormo.py` in `model_dir`) with the Auto classes.""" |
| cfg = get_class_from_dynamic_module("_configuration_kormo.KORMoConfig", model_dir) |
| base = get_class_from_dynamic_module("_modeling_kormo.KORMoModel", model_dir) |
| lm = get_class_from_dynamic_module("_modeling_kormo.KORMoForCausalLM", model_dir) |
| for fn in (lambda: AutoConfig.register("kormo", cfg), lambda: AutoModel.register(cfg, base), lambda: AutoModelForCausalLM.register(cfg, lm)): |
| try: |
| fn() |
| except ValueError: |
| pass |
|
|
|
|
| def patch_vision(model): |
| """The model was trained on SigLIP features *after* post_layernorm (transformers 4.57 behaviour). |
| In transformers 5, hidden_states[-1] is the value *before* post_layernorm, so replace it with last_hidden_state. |
| Call once after from_pretrained; without it the model still runs but misreads screenshots.""" |
| vt = model.model.vision_tower if hasattr(model, "model") else model.vision_tower |
|
|
| def hook(mod, args, out): |
| if getattr(out, "hidden_states", None): |
| out.hidden_states = tuple(out.hidden_states[:-1]) + (out.last_hidden_state,) |
| return out |
|
|
| vt.register_forward_hook(hook) |
| return model |
|
|
