KSA-V2-1.9B-base / indexer_step.py
OpenOneRec's picture
Release KSA-V2-1.9B-base weights and inference runtime
03c3dff verified
Raw History Blame Contribute Delete
1.26 kB
from contextvars import ContextVar
from functools import wraps
_INDEXER_STEP = ContextVar("indexer_step", default=None)
def shared_indexer_step(forward):
@wraps(forward)
def scoped(*args, **kwargs):
token = _INDEXER_STEP.set({"positions": {}, "bounds": {}})
try:
return forward(*args, **kwargs)
finally:
_INDEXER_STEP.reset(token)
return scoped
def shared_indexer_positions(source, kind, factory, *layout):
scope = _INDEXER_STEP.get()
if scope is None:
return factory()
key = (id(source), kind, *layout)
if key not in scope["positions"]:
# Retain the source, so tensor identity cannot be recycled in this step.
scope["positions"][key] = (source, factory())
return scope["positions"][key][1]
def _validated_position_bounds(position_ids):
scope = _INDEXER_STEP.get()
key = id(position_ids)
if scope is not None and key in scope["bounds"]:
return scope["bounds"][key][1]
low = int(position_ids.min())
if low < 0:
raise ValueError("indexer position IDs must be non-negative")
bounds = (low, int(position_ids.max()))
if scope is not None:
scope["bounds"][key] = (position_ids, bounds)
return bounds