How difra100/pt-rag_iclr works
PT-RAG predicts gene-perturbation responses in single-cell data. Its main workflow is to prepare perturbation data, train a state-transition model—optionally with differentiable, cell-type-aware retrieval—and run inference to generate and evaluate predictions. The repository also exposes a distinct embedding workflow for fitting or transforming cell embeddings and querying a LanceDB vector store. The README documents these capabilities; sampled source confirms CLI dispatch, embedding query behavior, data/model components, and prediction-time RAG handling, but does not expose every training or preprocessing call.
Main entry points: _predict.py Prediction Runner · state_transition.py State Transition Models · _tx TX Commands · _emb Embedding Commands
Components 22
Command Interface
- State CLI
src/state/__main__.py - TX Commands
src/state/_cli/_tx - Embedding Commands
src/state/_cli/_emb - Training Runner
src/state/_cli/_tx/_train.py - Prediction Runner
src/state/_cli/_tx/_predict.py - Data Preparation
src/state/_cli/_tx/_preprocess_train.py
Perturbation Modeling
- Perturbation Dataset
src/state/tx/data/dataset/scgpt_perturbation_dataset.py - State Transition Models
src/state/tx/models/state_transition.py - Model Variants
src/state/tx/models - scGPT Components
src/state/tx/models/scgpt - Differentiable Retrieval
src/state/tx/models/state_transition.py - Predicted Responses
Embedding Workflow
- Embedding Training
src/state/emb/train/trainer.py - Embedding Model
src/state/emb/nn/model.py - Embedding Data Loader
src/state/emb/data/loader.py - Embedding Inference
src/state/emb/inference.py - Vector Query
src/state/_cli/_emb/_query.py
Evaluation And Storage
- State Vector DB
src/state/emb/vectordb.py - LanceDB
- GenGeneEval
- Evaluation Metrics
src/state/_cli/_tx/_predict.py
Other components
- Researcher
Connections 28
- Researcher → State CLI: invokes
- State CLI → TX Commands: dispatches
- State CLI → Embedding Commands: dispatches
- TX Commands → Training Runner: selects train
- TX Commands → Prediction Runner: selects predict
- TX Commands → Data Preparation: selects preparation
- Data Preparation → Perturbation Dataset: prepares inputs
- Training Runner → Perturbation Dataset: loads batches
- Training Runner → State Transition Models: trains
- State Transition Models → Differentiable Retrieval: learns retrieval
- State Transition Models → Model Variants: uses configured model
- State Transition Models → scGPT Components: uses components
- Prediction Runner → Perturbation Dataset: sets up test data
- Prediction Runner → State Transition Models: loads checkpoint
- Prediction Runner → Predicted Responses: writes responses
- Prediction Runner → GenGeneEval: evaluates with
- GenGeneEval → Evaluation Metrics: returns metrics
- Embedding Commands → Embedding Training: selects fit
- Embedding Training → Embedding Data Loader: loads data
- Embedding Training → Embedding Model: trains
- Embedding Commands → Embedding Inference: selects transform
- Embedding Inference → Embedding Data Loader: loads cells
- Embedding Inference → Embedding Model: embeds cells
- Embedding Inference → State Vector DB: updates index
- State Vector DB → LanceDB: persists vectors
- Embedding Commands → Vector Query: selects query
- Vector Query → State Vector DB: searches vectors
- Vector Query → LanceDB: queries table