PECOS - Predictions for Enormous and Correlated Output Spaces
PECOS is a versatile and modular machine learning (ML) framework for fast learning and inference on problems with large output spaces, such as extreme multi-label ranking (XMR) and large-scale retrieval. PECOS' design is intentionally agnostic to the specific nature of the inputs and outputs as it is envisioned to be a general-purpose framework for multiple distinct applications.
Given an input, PECOS identifies a small set (10-100) of relevant outputs from amongst an extremely large (~100MM) candidate set and ranks these outputs in terms of relevance.
Features
Extreme Multi-label Ranking and Classification
-
X-Linear (
pecos.xmc.xlinear): recursive linear models learning to traverse an input from the root of a hierarchical label tree to a few leaf node clusters, and return top-k relevant labels within the clusters as predictions. See more details in the PECOS paper (Yu et al., 2020).- fast real-time inference in C++
- can handle 100MM output space
-
XR-Transformer (
pecos.xmc.xtransformer): Transformer based XMC framework that fine-tunes pre-trained transformers recursively on multi-resolution objectives. It can be used to generate top-k relevant labels for a given instance or simply as a fine-tuning engine for task aware embeddings. See technical details in XR-Transformer paper (Zhang et al., 2021).- easy to extend with many pre-trained Transformer models from huggingface transformers.
- establishes the State-of-the-art on public XMC benchmarks.
-
ANN Search with HNSW (
pecos.ann.hnsw): a PECOS Approximated Nearest Neighbor (ANN) search module that implements the Hierarchical Navigable Small World Graphs (HNSW) algorithm (Malkov et al., TPAMI 2018).- Supports both sparse and dense input features
- SIMD optimization for both dense/sparse distance computation
- Supports thread-safe graph construction in parallel on multi-core shared memory machines
- Supports thread-safe Searchers to do inference in parallel, which reduces inference overhead
Requirements and Installation
- Python (>=3.6)
- Pip (>=19.3)
See other dependencies in setup.py You should install PECOS in a virtual environment. If you're unfamiliar with Python virtual environments, check out the user guide.
Supporting Platforms
- Ubuntu 18.04 and 20.04
- Amazon Linux 2
Installation from Wheel
PECOS can be installed using pip as follows:
python3 -m pip install libpecos
Installation from Source
Prerequisite builder tools
- For Ubuntu (18.04, 20.04):
sudo apt-get update && sudo apt-get install -y build-essential git python3 python3-distutils python3-venv
- For Amazon Linux 2 Image:
sudo yum -y install python3 python3-devel python3-distutils python3-venv && sudo yum -y install groupinstall 'Development Tools'
One needs to install at least one BLAS library to compile PECOS, e.g. OpenBLAS:
- For Ubuntu (18.04, 20.04):
sudo apt-get install -y libopenblas-dev
- For Amazon Linux 2 Image and AMI:
sudo amazon-linux-extras install epel -y
sudo yum install openblas-devel -y
Install and develop locally
git clone https://github.com/amzn/pecos
cd pecos
python3 -m pip install --editable ./
Quick Tour
To have a glimpse of how PECOS works, here is a quick tour of using PECOS API for the XMR problem.
Toy Example
The eXtreme Multi-label Ranking (XMR) problem is defined by two matrices
- instance-to-feature matrix
X, of shapeN by DinSciPy CSR format - instance-to-label matrix
Y, of shapeN by LinSciPy CSR format
Some toy data matrices are available in the tst-data folder.
PECOS constructs a hierarchical label tree and learns linear models recursively (e.g., XR-Linear):
>>> from pecos.xmc.xlinear.model import XLinearModel
>>> from pecos.xmc import Indexer, LabelEmbeddingFactory
# Build hierarchical label tree and train a XR-Linear model
>>> label_feat = LabelEmbeddingFactory.create(Y, X)
>>> cluster_chain = Indexer.gen(label_feat)
>>> model = XLinearModel.train(X, Y, C=cluster_chain)
>>> model.save("./save-models")
After learning the model, we do prediction and evaluation
>>> from pecos.utils import smat_util
>>> Yt_pred = model.predict(Xt)
# print precision and recall at k=10
>>> print(smat_util.Metrics.generate(Yt, Yt_pred))
PECOS also offers optimized C++ implementation for fast real-time inference
>>> model = XLinearModel.load("./save-models", is_predict_only=True)
>>> for i in range(X_tst.shape[0]):
>>> y_tst_pred = model.predict(X_tst[i], threads=1)
Citation
If you find PECOS useful, please consider citing the following paper:
Some papers from our group using PECOS:
-
Accelerating Inference for Sparse Extreme Multi-Label Ranking Trees (Etter et al., 2021) [bib]
-
Enabling Efficiency-Precision Trade-offs for Label Trees in Extreme Classification (Baharav et al., 2021) [bib]
-
Fast Multi-Resolution Transformer Fine-tuning for Extreme Multi-label Text Classification (Zhang et al., NeurIPS 2021) [bib]
-
Label Disentanglement in Partition-based Extreme Multilabel Classification (Liu et al., NeurIPS 2021) [bib]
-
Extreme Multi-label Learning for Semantic Matching in Product Search (Chang et al., KDD 2021) [bib]
-
Session-Aware Query Auto-completion using Extreme Multi-label Ranking (Yadav et al., KDD 2021) [bib]
-
Top-k eXtreme Contextual Bandits with Arm Hierarchy (Sen et al., ICML 2021) [bib]
-
Taming pretrained transformers for extreme multi-label text classification (Chang et al., KDD 2020) [bib]
License
Copyright (2021) Amazon.com, Inc.
Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License.
203 Jan 3, 2023
34 Nov 21, 2022
21 Aug 30, 2022
62 Dec 27, 2022
1 Dec 29, 2021
1 Jan 10, 2022
69 Dec 26, 2022
75 Dec 15, 2022
77 Dec 22, 2022
76 Nov 23, 2022
0 Jan 23, 2022
58 May 3, 2022
75 Dec 30, 2022
130 Dec 13, 2022
1.8k Jan 5, 2023
69 Dec 17, 2022
74 Dec 3, 2022
40 Nov 18, 2022
614 Dec 31, 2022