Views
No views yet
res_kmeans.py - Model definitiontrain_res_kmeans.py - Training scriptinfer_res_kmeans.py - Inference scriptpip install torch numpy pandas pyarrow faiss tqdm1python train_res_kmeans.py \
2 --data_path ./data/embeddings.parquet \
3 --model_path ./checkpoints \
4 --n_layers 3 \
5 --codebook_size 8192 \
6 --dim 4096--data_path: Path to parquet file(s) with embedding column--model_path: Directory to save the model--n_layers: Number of residual layers (default: 3)--codebook_size: Size of each codebook (default: 8192)--dim: Embedding dimension (default: 4096)--seed: Random seed (default: 42)1python infer_res_kmeans.py \
2 --model_path ./checkpoints/model.pt \
3 --emb_path ./data/embeddings.parquet \
4 --output_path ./output/codes.parquet--model_path: Path to trained model checkpoint--emb_path: Path to parquet file with pid and embedding columns--output_path: Output path (default: {emb_path}_codes.parquet)--batch_size: Inference batch size (default: 10000)--device: Device to use (default: cuda if available)--n_layers: Number of layers to use (default: all)pid, embeddingpid, codes