Skip to content

Improve gene vocabulary loading and torchtext compatibility - #354

Open
Evanescence0515 wants to merge 1 commit into
bowang-lab:mainfrom
Evanescence0515:fix/gene-vocab-loading-compat
Open

Improve gene vocabulary loading and torchtext compatibility#354
Evanescence0515 wants to merge 1 commit into
bowang-lab:mainfrom
Evanescence0515:fix/gene-vocab-loading-compat

Conversation

@Evanescence0515

Copy link
Copy Markdown

Summary

This PR improves gene vocabulary loading performance and removes the direct torchtext dependency from the integration fine-tuning example.

It addresses the slow or apparently hanging vocabulary construction reported in #267 by replacing repeated token insertion with validated bulk initialization.

It also removes version-sensitive imports of torchtext.vocab.Vocab and the private torchtext._torchtext.Vocab implementation from examples/finetune_integration.py, related to the torchtext compatibility problem reported in #352.

Related issues: #267, #352.

Problem

GeneVocab.from_dict() currently inserts tokens one at a time through insert_token().

With the pure-Python BuiltinVocab backend, each insertion rebuilds the complete token-to-index mapping. Loading a vocabulary with approximately 60,000 tokens therefore has quadratic time complexity and can take more than one minute.

In addition, examples/finetune_integration.py directly imports torchtext.vocab.Vocab and the private torchtext._torchtext.Vocab implementation. This can fail across different torchtext versions.

Changes

  • Validate that vocabulary indices are integers, unique, consecutive, and start at zero.
  • Initialize ordered vocabulary tokens in bulk through _init_from_tokens() instead of repeatedly calling insert_token().
  • Preserve the exact token-to-index mapping and default-token behavior.
  • Remove direct torchtext imports from examples/finetune_integration.py.
  • Use GeneVocab.from_dict() when constructing a vocabulary without a pretrained model, preserving the original token order.
  • Add regression tests for unordered mappings and invalid indices.

Performance

Using the built-in vocabulary backend, a vocabulary containing 60,697 tokens was constructed in approximately:

0.0181 seconds

Tests

python -m pytest tests/test_tokenizer.py -q
15 passed, 3 skipped

The skipped tests require the optional torchtext dependency.
Additional formatting, syntax, and patch checks passed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant