[refactor] expose features/feature_groups as @property across models and wrappers#519
Conversation
…classes Add `features` / `feature_groups` `@property` accessors to the four classes that own the underlying `_features` / `_feature_groups` fields. Each property is a default read of the corresponding underscore field -- non-functional today, but gives subclasses a seam to override the surface per view (training vs export) without mutating the underscore fields. Touched: * `BaseModel` (tzrec/models/model.py): covers every ranking / matching / multi-task model and its `MatchModel` subclass. * `MatchTower` (tzrec/models/match_model.py): the base for `DSSMTower`, `DATTower`, `MINDUserTower`, `MINDItemTower`. Inherits from `BaseModule`, which does not own the underscore fields, so the property must be defined here for `TowerWrapper`'s property-based read (next commits) to succeed. * `MatchTowerWoEG` (tzrec/models/match_model.py): the base for `HSTUUserTower` / `HSTUMatchItemTower`. Inherits from `nn.Module` for the same reason. * `TDMEmbedding` (tzrec/models/tdm.py): an `nn.Module` exported separately by the TDM pipeline; owns its own `_features` / `_feature_groups` from the parent `EmbeddingGroup`. Pure additions -- no caller is rewired yet, so the existing underscore-field reads on these classes and their wrappers continue to work. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
… property API Switch the five call sites that read `model._features` / `model._feature_groups` (where `model` is a wrapper: `ScriptWrapper(...)` for non-match, `InferWrapper(TowerWrapper(...))` or `InferWrapper(TowerWoEGWrapper(...))` for match) to the property API introduced in the previous commit. Sites migrated: * `tzrec/utils/export_util.py::export_model_normal` — `features` argument to `create_dataloader` (line 164). * `tzrec/utils/export_util.py::export_model_aot` setup — second dataloader site (line 714). * `tzrec/utils/export_util.py::split_model` — `_compute_seq_share_groups` call site reads both `features` and `feature_groups` (lines 1163-1164). * `tzrec/acc/aot_utils.py::export_unified_model_aot` — sequence feature name set (line 470) and `_build_dynamic_shapes` call (lines 476-477). * `tzrec/acc/aot_utils.py::_build_dynamic_shapes` docstring — `model._features` / `model._feature_groups` references updated to `model.features` / `model.feature_groups` to match the new API. Today the wrappers still expose the underscore-field snapshot, so both forms work; this commit threads the property reads through the export pipeline so the next commit (which drops the underscore-field snapshots from the wrappers) is a no-op for these call sites. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…se properties `ScriptWrapper`, `TowerWrapper`, and `TowerWoEGWrapper` previously mirrored their inner module's `_features` / `_feature_groups` onto themselves at construction time. With every wrapped class now exposing those metadata fields via `@property` (previous commits), the wrappers can drop the snapshot and forward live reads instead. Concretely: * `ScriptWrapper.__init__`: no longer copies `model._features` / `model._feature_groups`. `DataParser` is constructed from `self.model.features`. The class exposes `features` / `feature_groups` as `@property` reads of `self.model.features` / `self.model.feature_groups`. * `TowerWrapper.__init__` / `TowerWoEGWrapper.__init__`: same treatment, reading the wrapped tower via `getattr(self, self._tower_name)`. * `TowerWoEGWrapper.__init__` additionally switches its `EmbeddingGroup` build from `module._features, module._feature_groups` to `module.features, module.feature_groups`, so the construction-time snapshot that EmbeddingGroup needs (it owns nn.Parameters) goes through the same property surface as everything else. The downstream call sites in `tzrec/utils/export_util.py` and `tzrec/acc/aot_utils.py` were migrated to the property API in the previous commit, so this is a no-op for them. No other consumer reads the wrappers' underscore fields directly. Non-functional for DSSM/MIND/DAT/TDM/HSTU: properties return the same `_features` / `_feature_groups` lists the wrappers used to snapshot, identical to the pre-refactor behaviour. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
| self._features = module._features | ||
| self._feature_groups = module._feature_groups | ||
| self._tower_name = tower_name | ||
| self._group_name = module._group_name |
There was a problem hiding this comment.
Inconsistent with the rest of the migration: this wrapper now reads features / feature_groups through the property surface but still reaches into module._group_name directly. The same "let subclasses override per-view" rationale applies to group_name — a MatchTowerWoEG subclass that wanted to expose a different group_name at export would have no seam.
Consider either (a) adding a group_name property on MatchTower / MatchTowerWoEG and reading through it here, or (b) at minimum a one-line comment noting _group_name is deliberately kept as a private contract between the tower and its wrapper. Symmetry with the rest of the PR would argue for (a).
Review summaryTightly-scoped, mechanical refactor. The shape is right: read-only Coverage: TorchScript: confirmed neither Security: no concerns. One noteworthy finding (inline at Nothing else blocking. |
…eature-groups-as-properties
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
2 participants
Summary
Add
features/feature_groups@propertyaccessors to the classes that own the underlying_features/_feature_groupslists (BaseModel,MatchTower,MatchTowerWoEG,TDMEmbedding); migrate wrappers (ScriptWrapper,TowerWrapper,TowerWoEGWrapper) and downstream export callers to read through that property surface so the wrappers no longer have to snapshot the inner module's metadata at construction time.Today every wrapper mirrors the inner module's
_features/_feature_groupsonto itself. That snapshot is fine for the static-view models we have, but it pins each wrapper to a single training-shape view. Exposing the lists as@propertys gives subclasses a seam to override the surface per view (training vs. export) without mutating the underscore fields the snapshot was reading from.Extracted from #518 to land independently. The HSTUMatch scalar export view in #518 will sit on top of this —
HSTUMatchItemToweroverrides the property to return a lazily-built scalar view when_is_inferenceis set.Design
_features/_feature_groupsunderscore fields stay as the source of truth onBaseModel,MatchTower,MatchTowerWoEG,TDMEmbedding. The new properties are default reads of those fields; existing subclasses pick them up unchanged.getattr(self, self._tower_name).featuresfor the per-tower wrappers;self.model.featuresforScriptWrapper).TowerWoEGWrapper.__init__switches itsEmbeddingGroupbuild tomodule.features/module.feature_groups. TheEmbeddingGroupownsnn.Parametersso it must remain a construction-time snapshot — but it now reads through the same property surface as everything else, so a subclass that overrides the property to a different view (e.g. HSTUMatch's scalar export view in [feat] HSTUMatch: scalar item-tower export view #518) gets the matchingEmbeddingGroup.tzrec/utils/export_util.pyandtzrec/acc/aot_utils.pymigrate frommodel._features/model._feature_groupsto the property API. Five sites total, all of which receive a wrapper (InferWrapper(...)orInferWrapper(TowerWrapper/TowerWoEGWrapper(...))).MatchTowerproperty — completeness fix vs #518#518 adds the property to
MatchTowerWoEGonly. Once wrapper snapshots are dropped,TowerWrapper.featuresreadsmodule.features— andMatchTowerinherits fromBaseModule, which has no such property. That would silently breakDSSMTower(dssm.py),DATTower(dat.py),MINDUserTower/MINDItemTower(mind.py) at export. This PR adds the missing property toMatchTowerso DSSM/DAT/MIND continue to work.Non-functional for DSSM/MIND/DAT/TDM/HSTU
The properties return the same
_features/_feature_groupslists the wrappers used to snapshot. No view changes, no model changes — purely a name-routing refactor that unlocks per-view overrides for future PRs (notably #518's HSTUMatch scalar export).Test plan
python -m unittest tzrec.models.match_model_test— green (4 tests, MatchModel + TrainWrapper)python -m unittest tzrec.models.rank_model_test— green (10 tests, RankModel/BaseModel surface)python -m unittest tzrec.models.tdm_test— green (3 tests, TDM + TDMEmbedding)python -m unittest tzrec.models.dssm_test— green (6 tests, MatchTower path)python -m unittest tzrec.models.dssm_v2_test— green (2 tests, MatchTowerWoEG path)python -m unittest tzrec.models.hstu_test— green (1 test, MatchTowerWoEG via HSTUMatchItemTower)python -m unittest tzrec.models.mind_test— green (6 tests, MatchTower path)python -m unittest tzrec.models.dat_test— green (6 tests, MatchTower path)pre-commit run --files <touched>— cleanTowerWrapper/TowerWoEGWrapper/ScriptWrapper).🤖 Generated with Claude Code