"""Fix SGLang qwen3_5_text loader: yield body-relative names.

Root cause: the wrapper stored the body as self.model, so the body's own
params_dict holds PREFIX-LESS names (layers.0..., embed_tokens, norm...).
But body_weights() forwarded names that never match those keys:
  - original code stripped only "model." -> "language_model.layers..." which
    the body's "model.language_model."->"model." rule cannot map;
  - checkpoint names are "model.language_model.*" (VL-structured) or
    "model.*" (plain), neither of which resolves.
Almost every weight missed params_dict ("not found" warnings) and the model
ran on uninitialized params (constant token-0 output).

Fix: strip "model." AND an optional "language_model." segment, yielding
exactly the body-relative names for both checkpoint layouts. lm_head keeps
its dedicated branch. Idempotent.
"""
import sys

P = "/sgl-workspace/sglang/python/sglang/srt/models/qwen3_5_text.py"
src = open(P).read()
if "# PREFIX-STRIP-PATCH-V2" in src:
    print("already patched (v2)")
    sys.exit(0)

v1 = """                if name.startswith(_MODEL_PREFIX):
                    # NO-STRIP-PATCH: keep full names; the body maps
                    # model.language_model.* -> model.* itself. Stripping
                    # breaks every lookup (params stay uninitialized).
                    yield name, loaded_weight"""
v1_orig = """                if name.startswith(_MODEL_PREFIX):
                    yield name[len(_MODEL_PREFIX) :], loaded_weight"""
if v1 in src:
    src = src.replace(v1, "PLACEHOLDER_WILL_REPLACE_BELOW")
    print("removed v1 patch")
elif v1_orig not in src:
    print("pattern not found - SGLang version differs, aborting")
    sys.exit(1)

fixed = """                # PREFIX-STRIP-PATCH-V2: the body stores params WITHOUT any
                # prefix (layers.0..., embed_tokens, ...), so strip "model."
                # plus an optional "language_model." segment. Covers both
                # "model.layers.*" and "model.language_model.*" checkpoints.
                if name.startswith("model.language_model."):
                    yield name[len("model.language_model."):], loaded_weight
                elif name.startswith(_MODEL_PREFIX):
                    yield name[len(_MODEL_PREFIX):], loaded_weight"""
if "PLACEHOLDER_WILL_REPLACE_BELOW" in src:
    src = src.replace("PLACEHOLDER_WILL_REPLACE_BELOW", fixed)
else:
    src = src.replace(v1_orig, fixed)
open(P, "w").write(src)
print("patched OK (v2)")