feat(context): Add TokenBudget and pluggable tokenizers - #6824
Open
patelchaitany wants to merge 1 commit into
Open
feat(context): Add TokenBudget and pluggable tokenizers#6824patelchaitany wants to merge 1 commit into
patelchaitany wants to merge 1 commit into
Conversation
Create the feast.context module with the token accounting that prompt assembly in an OnDemandFeatureView builds on. Tokenizers resolve the way online store types do in repo_config: a built-in name maps to a class path in TOKENIZER_CLASS_FOR_TYPE, and anything else is itself the path of a Tokenizer subclass with a no-argument constructor, loaded through import_class. Adding a tokenizer is therefore what adding a vector store is - write the class, pass its path, no registration call and no mutable global state. Feast ships cl100k_base, o200k_base and a dependency-free character estimate. TokenBudget is a frozen dataclass holding a resolved Tokenizer: consume() returns a new budget rather than mutating, and try_consume() returns None instead of raising, which is the primitive priority_select() will use to fill a budget greedily. TokenBudget.of() resolves a name, class path or instance and can refuse the fallback; is_approximate reports whether counts are exact. Tokenizers compare by value, so budgets built from the same name are interchangeable as dict keys and set members. When tiktoken cannot be loaded, get_tokenizer() degrades to the character estimate with a logged warning; fallback=False makes it fatal for paths that need exact counts. TiktokenTokenizer counts with disallowed_special=(), so a feature value containing a literal "<|endoftext|>" is treated as ordinary text instead of raising. tiktoken is not declared as a dependency yet, so the tests that assert exact encodings skip when it is absent. Part of RHOAIENG-80116 (PR 1/9). Signed-off-by: Chaitany Patel <patelchaitany93@gmail.com>
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## master #6824 +/- ##
=======================================
Coverage 47.08% 47.08%
=======================================
Files 419 419
Lines 51877 51877
Branches 7525 7525
=======================================
Hits 24428 24428
Misses 25700 25700
Partials 1749 1749
*This pull request uses carry forward flags. Click here to find out more. Continue to review full report in Codecov by Harness.
🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this PR does / why we need it:
Adds
feast.context, a new module holding the token accounting that prompt assembly in an OnDemandFeatureView will build on. This is the first of a planned series; nothing in Feast imports it yet, so the change is additive and inert.Tokenizers are pluggable the same way online stores and vector stores already are.
TOKENIZER_CLASS_FOR_TYPEmaps a built-in name to a class path, exactly asONLINE_STORE_CLASS_FOR_TYPEdoes inrepo_config.py, and anything that is not a built-in name is itself taken to be the fully-qualified path of aTokenizersubclass, loaded throughfeast.importer.import_class. There is no registry and no mutable global state — adding a tokenizer means writing the class and passing its path:Feast ships
cl100k_baseando200k_base(tiktoken) plus a dependency-free character estimate.TokenBudgetis a frozen dataclass.consume()returns a new budget instead of mutating, so a budget is safe to share across threads and to reuse between assemblies;try_consume()returnsNonerather than raising, which is the primitive that greedy section selection will use later. Tokenizers compare by value, so budgets built from the same tokenizer name are interchangeable as dict keys and set members.Degradation is explicit rather than silent. When tiktoken cannot be loaded,
get_tokenizer()falls back to the character estimate and logs a warning;TokenBudget.of(..., fallback=False)makes that fatal for paths needing exact counts, andis_approximatelets a caller check without matching on the tokenizer name.TiktokenTokenizercounts withdisallowed_special=(), so a feature value containing a literal<|endoftext|>is treated as ordinary text instead of raising.Note on the dependency: tiktoken is not declared here, so the tests asserting exact encodings skip when it is absent. Happy to add a
feast[llm]extra in this PR or a follow-up, whichever reviewers prefer.Which issue(s) this PR fixes:
No GitHub issue. Tracked internally as RHOAIENG-80116.
Checks
git commit -s)Testing Strategy
sdk/python/tests/unit/context/— 61 tests with tiktoken installed, 51 passing and 10 skipped without it. Covers name and class-path resolution, rejection of a class that is merely namedTokenizer, the fallback in both directions, budget arithmetic and immutability, and exact counts for both shipped encodings.Misc
New module, no existing behaviour touched.
ruff check,ruff format --checkandmypyare clean across the eight files.