├── .gitignore ├── LICENSE ├── README.md ├── __init__.py ├── configs ├── __init__.py ├── callbacks │ ├── default.yaml │ ├── early_stopping.yaml │ ├── local_pickle_writer.yaml │ ├── model_checkpoint.yaml │ ├── model_summary.yaml │ └── rich_progress_bar.yaml ├── experiment │ ├── rkmeans_inference_flat.yaml │ ├── rkmeans_train_flat.yaml │ ├── rqvae_train_flat.yaml │ ├── rvq_train_flat.yaml │ ├── sem_embeds_inference_flat.yaml │ ├── tiger_inference_flat.yaml │ └── tiger_train_flat.yaml ├── extras │ └── default.yaml ├── hydra │ └── default.yaml ├── inference.yaml ├── local │ └── .gitkeep ├── logger │ └── csv.yaml ├── paths │ └── default.yaml ├── train.yaml └── trainer │ ├── ddp.yaml │ └── default.yaml ├── logs └── .gitkeep ├── notices.txt ├── outputs └── .gitkeep ├── requirements.txt └── src ├── __init__.py ├── components ├── __init__.py ├── clustering_initializers.py ├── distance_functions.py ├── eval_metrics.py ├── loss_functions.py ├── network_blocks │ ├── __init__.py │ └── hf_language_model.py ├── optimizer.py ├── quantization_strategies.py ├── scheduler.py └── training_loop_functions.py ├── data └── loading │ ├── __init__.py │ ├── components │ ├── __init__.py │ ├── collate_functions.py │ ├── custom_dataloader.py │ ├── dataloading.py │ ├── interfaces.py │ ├── iterators.py │ ├── label_function.py │ └── pre_processing.py │ ├── datamodules │ ├── __init__.py │ └── sequence_datamodule.py │ └── utils.py ├── inference.py ├── models ├── __init__.py ├── components │ ├── __init__.py │ ├── interfaces.py │ └── network_blocks │ │ ├── aggregation_strategy.py │ │ ├── embedding_aggregator.py │ │ ├── mlp.py │ │ └── normalize_layer.py └── modules │ ├── __init__.py │ ├── base_module.py │ ├── clustering │ ├── base_clustering_module.py │ └── mini_batch_kmeans.py │ ├── huggingface │ ├── __init__.py │ └── transformer_base_module.py │ └── semantic_id │ ├── __init__.py │ └── tiger_generation_model.py ├── modules ├── __init__.py ├── clustering │ ├── __init__.py │ ├── residual_quantization.py │ └── vector_quantization.py └── semantic_embedding_inference_module.py ├── train.py └── utils ├── __init__.py ├── custom_hydra_resolvers.py ├── decorators.py ├── file_utils.py ├── inference_utils.py ├── instantiators.py ├── launcher_utils.py ├── logging_utils.py ├── masking_utils.py ├── pylogger.py ├── restart_job.py ├── restart_job_utils.py ├── rich_utils.py ├── tensor_utils.py └── utils.py /.gitignore: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/.gitignore -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/README.md -------------------------------------------------------------------------------- /__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /configs/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/__init__.py -------------------------------------------------------------------------------- /configs/callbacks/default.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/callbacks/default.yaml -------------------------------------------------------------------------------- /configs/callbacks/early_stopping.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/callbacks/early_stopping.yaml -------------------------------------------------------------------------------- /configs/callbacks/local_pickle_writer.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/callbacks/local_pickle_writer.yaml -------------------------------------------------------------------------------- /configs/callbacks/model_checkpoint.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/callbacks/model_checkpoint.yaml -------------------------------------------------------------------------------- /configs/callbacks/model_summary.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/callbacks/model_summary.yaml -------------------------------------------------------------------------------- /configs/callbacks/rich_progress_bar.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/callbacks/rich_progress_bar.yaml -------------------------------------------------------------------------------- /configs/experiment/rkmeans_inference_flat.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/experiment/rkmeans_inference_flat.yaml -------------------------------------------------------------------------------- /configs/experiment/rkmeans_train_flat.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/experiment/rkmeans_train_flat.yaml -------------------------------------------------------------------------------- /configs/experiment/rqvae_train_flat.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/experiment/rqvae_train_flat.yaml -------------------------------------------------------------------------------- /configs/experiment/rvq_train_flat.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/experiment/rvq_train_flat.yaml -------------------------------------------------------------------------------- /configs/experiment/sem_embeds_inference_flat.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/experiment/sem_embeds_inference_flat.yaml -------------------------------------------------------------------------------- /configs/experiment/tiger_inference_flat.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/experiment/tiger_inference_flat.yaml -------------------------------------------------------------------------------- /configs/experiment/tiger_train_flat.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/experiment/tiger_train_flat.yaml -------------------------------------------------------------------------------- /configs/extras/default.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/extras/default.yaml -------------------------------------------------------------------------------- /configs/hydra/default.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/hydra/default.yaml -------------------------------------------------------------------------------- /configs/inference.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/inference.yaml -------------------------------------------------------------------------------- /configs/local/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /configs/logger/csv.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/logger/csv.yaml -------------------------------------------------------------------------------- /configs/paths/default.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/paths/default.yaml -------------------------------------------------------------------------------- /configs/train.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/train.yaml -------------------------------------------------------------------------------- /configs/trainer/ddp.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/trainer/ddp.yaml -------------------------------------------------------------------------------- /configs/trainer/default.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/configs/trainer/default.yaml -------------------------------------------------------------------------------- /logs/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /notices.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/notices.txt -------------------------------------------------------------------------------- /outputs/.gitkeep: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/requirements.txt -------------------------------------------------------------------------------- /src/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/components/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/components/clustering_initializers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/components/clustering_initializers.py -------------------------------------------------------------------------------- /src/components/distance_functions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/components/distance_functions.py -------------------------------------------------------------------------------- /src/components/eval_metrics.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/components/eval_metrics.py -------------------------------------------------------------------------------- /src/components/loss_functions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/components/loss_functions.py -------------------------------------------------------------------------------- /src/components/network_blocks/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/components/network_blocks/hf_language_model.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/components/network_blocks/hf_language_model.py -------------------------------------------------------------------------------- /src/components/optimizer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/components/optimizer.py -------------------------------------------------------------------------------- /src/components/quantization_strategies.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/components/quantization_strategies.py -------------------------------------------------------------------------------- /src/components/scheduler.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/components/scheduler.py -------------------------------------------------------------------------------- /src/components/training_loop_functions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/components/training_loop_functions.py -------------------------------------------------------------------------------- /src/data/loading/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/data/loading/components/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/data/loading/components/collate_functions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/data/loading/components/collate_functions.py -------------------------------------------------------------------------------- /src/data/loading/components/custom_dataloader.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/data/loading/components/custom_dataloader.py -------------------------------------------------------------------------------- /src/data/loading/components/dataloading.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/data/loading/components/dataloading.py -------------------------------------------------------------------------------- /src/data/loading/components/interfaces.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/data/loading/components/interfaces.py -------------------------------------------------------------------------------- /src/data/loading/components/iterators.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/data/loading/components/iterators.py -------------------------------------------------------------------------------- /src/data/loading/components/label_function.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/data/loading/components/label_function.py -------------------------------------------------------------------------------- /src/data/loading/components/pre_processing.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/data/loading/components/pre_processing.py -------------------------------------------------------------------------------- /src/data/loading/datamodules/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/data/loading/datamodules/sequence_datamodule.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/data/loading/datamodules/sequence_datamodule.py -------------------------------------------------------------------------------- /src/data/loading/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/data/loading/utils.py -------------------------------------------------------------------------------- /src/inference.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/inference.py -------------------------------------------------------------------------------- /src/models/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/models/components/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/models/components/interfaces.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/models/components/interfaces.py -------------------------------------------------------------------------------- /src/models/components/network_blocks/aggregation_strategy.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/models/components/network_blocks/aggregation_strategy.py -------------------------------------------------------------------------------- /src/models/components/network_blocks/embedding_aggregator.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/models/components/network_blocks/embedding_aggregator.py -------------------------------------------------------------------------------- /src/models/components/network_blocks/mlp.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/models/components/network_blocks/mlp.py -------------------------------------------------------------------------------- /src/models/components/network_blocks/normalize_layer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/models/components/network_blocks/normalize_layer.py -------------------------------------------------------------------------------- /src/models/modules/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/models/modules/base_module.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/models/modules/base_module.py -------------------------------------------------------------------------------- /src/models/modules/clustering/base_clustering_module.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/models/modules/clustering/base_clustering_module.py -------------------------------------------------------------------------------- /src/models/modules/clustering/mini_batch_kmeans.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/models/modules/clustering/mini_batch_kmeans.py -------------------------------------------------------------------------------- /src/models/modules/huggingface/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/models/modules/huggingface/transformer_base_module.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/models/modules/huggingface/transformer_base_module.py -------------------------------------------------------------------------------- /src/models/modules/semantic_id/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/models/modules/semantic_id/tiger_generation_model.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/models/modules/semantic_id/tiger_generation_model.py -------------------------------------------------------------------------------- /src/modules/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/modules/clustering/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /src/modules/clustering/residual_quantization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/modules/clustering/residual_quantization.py -------------------------------------------------------------------------------- /src/modules/clustering/vector_quantization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/modules/clustering/vector_quantization.py -------------------------------------------------------------------------------- /src/modules/semantic_embedding_inference_module.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/modules/semantic_embedding_inference_module.py -------------------------------------------------------------------------------- /src/train.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/train.py -------------------------------------------------------------------------------- /src/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/utils/__init__.py -------------------------------------------------------------------------------- /src/utils/custom_hydra_resolvers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/utils/custom_hydra_resolvers.py -------------------------------------------------------------------------------- /src/utils/decorators.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/utils/decorators.py -------------------------------------------------------------------------------- /src/utils/file_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/utils/file_utils.py -------------------------------------------------------------------------------- /src/utils/inference_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/utils/inference_utils.py -------------------------------------------------------------------------------- /src/utils/instantiators.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/utils/instantiators.py -------------------------------------------------------------------------------- /src/utils/launcher_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/utils/launcher_utils.py -------------------------------------------------------------------------------- /src/utils/logging_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/utils/logging_utils.py -------------------------------------------------------------------------------- /src/utils/masking_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/utils/masking_utils.py -------------------------------------------------------------------------------- /src/utils/pylogger.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/utils/pylogger.py -------------------------------------------------------------------------------- /src/utils/restart_job.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/utils/restart_job.py -------------------------------------------------------------------------------- /src/utils/restart_job_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/utils/restart_job_utils.py -------------------------------------------------------------------------------- /src/utils/rich_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/utils/rich_utils.py -------------------------------------------------------------------------------- /src/utils/tensor_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/utils/tensor_utils.py -------------------------------------------------------------------------------- /src/utils/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/snap-research/GRID/HEAD/src/utils/utils.py --------------------------------------------------------------------------------