41 lines
1.5 KiB
Python
41 lines
1.5 KiB
Python
import ast
|
|
from pathlib import Path
|
|
|
|
|
|
def _find_vision():
|
|
for p in [
|
|
Path(__file__).resolve().parent / "unsloth" / "models" / "vision.py",
|
|
Path(__file__).resolve().parents[1] / "unsloth" / "models" / "vision.py",
|
|
Path(
|
|
"/mnt/disks/unslothai/ubuntu/workspace_25/github_review/unsloth-pr-5053-staging-3/unsloth/models/vision.py"
|
|
),
|
|
]:
|
|
if p.exists():
|
|
return p
|
|
raise FileNotFoundError("vision.py not found")
|
|
|
|
|
|
def test_vision_fastbasemodel_from_pretrained_calls_helper():
|
|
"""FastBaseModel.from_pretrained must invoke _attach_bnb_multidevice_hooks
|
|
after the underlying model load so the inference hook path is reachable
|
|
via the base vision loader, not only via llama."""
|
|
src = _find_vision().read_text()
|
|
tree = ast.parse(src)
|
|
found = False
|
|
for cls in ast.walk(tree):
|
|
if not (isinstance(cls, ast.ClassDef) and cls.name == "FastBaseModel"):
|
|
continue
|
|
for fn in ast.walk(cls):
|
|
if not (isinstance(fn, ast.FunctionDef) and fn.name == "from_pretrained"):
|
|
continue
|
|
for node in ast.walk(fn):
|
|
if (
|
|
isinstance(node, ast.Call)
|
|
and getattr(node.func, "id", None)
|
|
== "_attach_bnb_multidevice_hooks"
|
|
):
|
|
found = True
|
|
break
|
|
assert (
|
|
found
|
|
), "FastBaseModel.from_pretrained must call _attach_bnb_multidevice_hooks"
|