Temporarily rename "TFT" resources for grepability#620
Conversation
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. WalkthroughRenamed TemporalFusionTransformerDataset to TFTDataset and TemporalFusionTransformer to TFTModel across library, training script, and tests. Updated function signatures, imports, and references to use the new class names. Added a docstring to TFTDataset. No behavioral or logic changes. Changes
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
This stack of pull requests is managed by Graphite. Learn more about stacking. |
There was a problem hiding this comment.
Pull Request Overview
This PR temporarily renames Temporal Fusion Transformer (TFT) related classes and functions to use abbreviated "TFT" naming for better code searchability and grepability.
- Renames
TemporalFusionTransformerDatasettoTFTDatasetacross test and application files - Renames
TemporalFusionTransformermodel class toTFTModel - Updates all corresponding function names and import statements to use the new abbreviated naming
Reviewed Changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| libraries/python/tests/test_tft_dataset.py | Updates test imports and class instantiations to use TFTDataset instead of TemporalFusionTransformerDataset |
| libraries/python/src/internal/tft_model.py | Renames main model class from TemporalFusionTransformer to TFTModel |
| libraries/python/src/internal/tft_dataset.py | Renames dataset class from TemporalFusionTransformerDataset to TFTDataset and adds docstring |
| applications/models/src/models/train_tft_model.py | Updates imports, type annotations, and function signatures to use new abbreviated class names |
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
Graphite Automations"Assign author to pull request" took an action on this PR • (08/22/25)1 assignee was added to this PR based on John Forstmeier's automation. |
… of github.com:pocketsizefund/pocketsizefund into 08-22-temporarily_rename_tft_resources_for_grepability
There was a problem hiding this comment.
Actionable comments posted: 0
♻️ Duplicate comments (1)
libraries/python/src/internal/tft_dataset.py (1)
28-30: Docstring is too terse; expand with purpose, schema, and usage.This mirrors a prior review asking for a meaningful docstring. Brief text hurts grepability and onboarding. Recommend a concise but informative block that captures expected columns, derived features, and what batching returns.
Apply this diff to enrich the docstring:
class TFTDataset: - """Temporal fusion transformer dataset.""" + """ + Dataset for Temporal Fusion Transformer (TFT). + + Responsibilities: + - Validate and preprocess structured time series data (one row per ticker × timestamp). + - Derive categorical time features, scale continuous features, and encode static categories. + - Produce train/validate/predict batches consumed by TFTModel. + + Expected input columns (exact match): + - ticker, timestamp (ms since epoch), open_price, high_price, low_price, + close_price, volume, volume_weighted_average_price, sector, industry + + Notable attributes: + - continuous_columns, categorical_columns, static_categorical_columns + - mappings: dict for categorical encodings + - data: preprocessed Polars DataFrame + + Key methods: + - get_dimensions() -> shapes of encoder/decoder/static feature groups + - get_batches(data_type={"train","validate","predict"}, …) -> list[dict[str, Tensor]] + """
🧹 Nitpick comments (2)
libraries/python/src/internal/tft_dataset.py (1)
28-30: Introduce a temporary alias for backward compatibility (optional)We ran a search for
TemporalFusionTransformerDatasetacross the codebase and found no existing references. However, adding a one-line alias will protect any downstream consumers who may still be importing the old class name during the transition:+# Temporary alias for backward compatibility during the rename window. +TemporalFusionTransformerDataset = TFTDataset # noqa: N816No further verification is needed—this is purely an optional refactor to smooth migration.
libraries/python/tests/test_tft_dataset.py (1)
2-2: Test updates align with the rename; add a predict-mode invariant and consider parameterization.
- Imports and test names reflect TFTDataset. Good.
- In predict mode, batches shouldn’t include targets. Adding an assertion hardens the contract.
- Consider pytest parameterization to reduce duplication across similar dataset constructions. Nice-to-have.
Apply this small assertion in the predict test:
for batch in batches: assert "encoder_categorical_features" in batch assert "encoder_continuous_features" in batch assert "decoder_categorical_features" in batch assert "static_categorical_features" in batch + assert "targets" not in batch # predict mode should not include targetsAlso applies to: 5-5, 42-42, 48-48, 67-67, 79-79, 103-103
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (4)
applications/models/src/models/train_tft_model.py(6 hunks)libraries/python/src/internal/tft_dataset.py(1 hunks)libraries/python/src/internal/tft_model.py(1 hunks)libraries/python/tests/test_tft_dataset.py(5 hunks)
🧰 Additional context used
🧬 Code graph analysis (3)
libraries/python/src/internal/tft_dataset.py (1)
libraries/python/src/internal/dataset.py (1)
TemporalFusionTransformerDataset(28-321)
libraries/python/tests/test_tft_dataset.py (2)
libraries/python/src/internal/tft_dataset.py (1)
TFTDataset(28-323)libraries/python/src/internal/dataset.py (1)
TemporalFusionTransformerDataset(28-321)
applications/models/src/models/train_tft_model.py (2)
libraries/python/src/internal/tft_dataset.py (1)
TFTDataset(28-323)libraries/python/src/internal/tft_model.py (2)
Parameters(19-33)TFTModel(37-273)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
- GitHub Check: Run Python tests
🔇 Additional comments (2)
libraries/python/src/internal/tft_model.py (1)
36-37: Alias and docstring added for backward compatibilityThe alias and minimal docstring have been applied as suggested, and a search for remaining references to the old class name returned no matches, so no further churn is expected.
• File affected:
libraries/python/src/internal/tft_model.py(around lines 36–37)
• Verified remaining references withrg -nP '\bTemporalFusionTransformer\b(?!Dataset)'(no hits)# https://arxiv.org/pdf/1912.09363 class TFTModel: """Temporal Fusion Transformer (TFT) model (temporarily renamed for grepability).""" @@ self.parameters = get_parameters(self) # Backward-compat alias during the rename window. TemporalFusionTransformer = TFTModel # noqa: N816applications/models/src/models/train_tft_model.py (1)
10-11: Patch loss logging & confirm W&B run lifecycleThe renames to TFTDataset/TFTModel look correct. Please apply the loss‐logging fix and verify whether ending the W&B run inside
train_tft_model.pyis intended—otherwise, move thewandb.finish()call to after validation/save so all metrics log under one run.• Loss logging
- losses = model.train( + train_stats = model.train( inputs_list=batches, epoch_count=epoch_count, learning_rate=learning_rate, ) - for loss in losses: - wandb_run.log({"loss": loss}) + for loss in train_stats["losses"]: + wandb_run.log({"loss": loss})• Optional logger message tweaks for consistency
- logger.info("Training temporal fusion transformer model") + logger.info("Training TFT model") @@ - logger.info("Validating temporal fusion transformer model") + logger.info("Validating TFT model") @@ - logger.info("Saving temporal fusion transformer model") + logger.info("Saving TFT model")• W&B lifecycle (in applications/models/src/models/train_tft_model.py)
- Line 136:
wandb.finish() # close active run if it exists- Line 138:
wandb_run = wandb.init(
Confirm whether you want to close the run before validation/save (current behavior), or deferwandb.finish()until all steps complete.

Overview
Changes
Comments
Just a minor tweak I cut as a separate pull request for clarity. Eventually
internal/will just need submodules (e.g.tft/) but that's easy cleanup.Summary by CodeRabbit
Refactor
Documentation
Tests