├── .gitignore ├── LICENSE ├── README.md ├── data ├── Grocery_and_Gourmet_Food │ ├── Amazon.ipynb │ ├── dev.csv │ ├── item_meta.csv │ ├── test.csv │ └── train.csv ├── MIND_Large │ └── MIND-large.ipynb ├── MovieLens_1M │ └── MovieLens-1M.ipynb └── README.md ├── docs ├── Getting_Started.md ├── Main_Arguments.md ├── Supported_Models.md ├── _static │ ├── ctr.png │ ├── data_format.png │ ├── format_meta.png │ ├── format_test.png │ ├── logo.png │ ├── logo2.0.png │ ├── module.png │ ├── module_new.png │ ├── rerank.png │ └── topk.png ├── demo_scripts_results │ ├── CTR_MIND.sh │ ├── CTR_ML1M.sh │ ├── README.md │ ├── Rerank_MIND.sh │ ├── Rerank_ML1M.sh │ ├── Topk_Amazon.sh │ ├── Topk_MIND.sh │ └── Topk_ML1M.sh └── tutorials │ ├── CTR_Prediction.ipynb │ ├── Context_aware_Topk_Recommendation.ipynb │ ├── Impression_based_Ranking_Reranking.ipynb │ └── README.md ├── requirements.txt └── src ├── README.md ├── exp.py ├── helpers ├── BUIRRunner.py ├── BaseReader.py ├── BaseRunner.py ├── CTRRunner.py ├── ContextReader.py ├── ContextSeqReader.py ├── ImpressionContextReader.py ├── ImpressionReader.py ├── ImpressionRunner.py ├── ImpressionSeqReader.py ├── KDAReader.py ├── KGReader.py ├── SeqReader.py └── __init__.py ├── main.py ├── models ├── BaseContextModel.py ├── BaseImpressionModel.py ├── BaseModel.py ├── BaseRerankerModel.py ├── __init__.py ├── context │ ├── AFM.py │ ├── AutoInt.py │ ├── DCN.py │ ├── DCNv2.py │ ├── DeepFM.py │ ├── FM.py │ ├── FinalMLP.py │ ├── SAM.py │ ├── WideDeep.py │ ├── __init__.py │ └── xDeepFM.py ├── context_seq │ ├── CAN.py │ ├── DIEN.py │ ├── DIN.py │ ├── ETA.py │ ├── SDIM.py │ └── __init__.py ├── developing │ ├── CLRec.py │ ├── FourierTA.py │ ├── S3Rec.py │ ├── SRGNN.py │ └── __init__.py ├── general │ ├── BPRMF.py │ ├── BUIR.py │ ├── CFKG.py │ ├── DirectAU.py │ ├── LightGCN.py │ ├── NeuMF.py │ ├── POP.py │ └── __init__.py ├── reranker │ ├── MIR.py │ ├── PRM.py │ ├── SetRank.py │ └── __init__.py └── sequential │ ├── Caser.py │ ├── Chorus.py │ ├── ComiRec.py │ ├── ContraRec.py │ ├── FPMC.py │ ├── GRU4Rec.py │ ├── KDA.py │ ├── NARM.py │ ├── SASRec.py │ ├── SLRCPlus.py │ ├── TiMiRec.py │ ├── TiSASRec.py │ └── __init__.py └── utils ├── __init__.py ├── layers.py └── utils.py /.gitignore: -------------------------------------------------------------------------------- 1 | .ipynb_checkpoints/ 2 | *.xlsx 3 | .idea/ 4 | __pycache__/ 5 | .DS_Store 6 | *.pyc 7 | figure/ 8 | doc/ 9 | model/ 10 | search.py 11 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2020 Chenyang Wang 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /data/README.md: -------------------------------------------------------------------------------- 1 | # Dataset 2 | 3 | We include the public [Amazon dataset](http://jmcauley.ucsd.edu/data/amazon/links.html) (*Grocery_and_Gourmet_Food* category, 5-core version with metadata), [MIND-Large dataset](https://msnews.github.io/), and [MovieLens-1M](https://grouplens.org/datasets/movielens/) as our built-in datasets. 4 | The pre-processed version of Amazon dataset can be found in the `./Grocery_and_Gourmet_Food` dir, which supports Top-k recommendation tasks. 5 | You can also download and process MIND and MovieLens datasets for CTR prediction, Top-k recommendation, and re-ranking tasks in the corresponding notebooks in `MIND_Large` and `MovieLens_1M` datasets. 6 | 7 | Our framework can also work with other datasets easily. We describe the required file formats for each task and the format for context information below: 8 | 9 | ## Top-k Recommendation Task 10 | 11 | **train.csv** 12 | - Format: `user_id \t item_id \t time` 13 | - All ids **begin from 1** (0 is reserved for NaN), and the followings are the same. 14 | - Need to be sorted in **time-ascending order** when running sequential models. 15 | 16 | **test.csv & dev.csv** 17 | 18 | - Format: `user_id \t item_id \t time \t neg_items` 19 | - The last column is the list of negative items corresponding to each ground-truth item (should not include unseen item ids beyond the `item_id` column in train/dev/test sets). 20 | - The number of negative items need to be the same for a specific set, but it can be different between dev and test sets. 21 | - If there is no `neg_items` column, the evaluation will be performed over all the items by default. 22 | 23 | ## CTR Prediction Task 24 | 25 | **train.csv & test.csv & dev.csv** 26 | - Format: `user_id \t item_id \t time \t label` 27 | - Labels should be 0 or 1 to indicate the item is clicked or not. 28 | - Need to be sorted in **time-ascending order** when running sequential models. 29 | 30 | ## Impression-based Ranking/Reranking Task 31 | 32 | **train.csv & test.csv & dev.csv** 33 | - Format: `user_id \t item_id \t time \t label \t impression_id` 34 | - All interactions with the same impression id will be grouped as a candidate list for training and evaluations. 35 | - If there is no `impression_id` column, interactions will grouped by `time`. 36 | - Labels should be 0 or 1 to indicate the item is clicked or not. 37 | - Need to be sorted in **time-ascending order** when running sequential models. 38 | 39 | 40 | ## Context Information 41 | 42 | **item_meta.csv** (optional) 43 | 44 | - Format: `item_id \t i__ \t ... \t r_ \t ...` 45 | - Optional, only needed for context-aware models and some of the knowledge-aware models (CFKG, SLRC+, Chorus, KDA). 46 | - For context-aware models, an argument called `include_item_features` is used to control whether to use the item metadata or not. 47 | - `i__` is the attribute of an item, such as category, brand, price, and so on. The features should be numerical. The header should start with `i_` and the is set to `c` for categorical features and `f` for dense (float) features. 48 | - `r_` is the relations between items, and its value is a list of items (can be empty []). Assume `item_id` is `i`, if `j` appears in `r_`, then `(i, relation, j)` holds in the knowledge graph. Note that the corresponding header here must start with "r_" to be distinguished from attributes. 49 | 50 | **user_meta.csv** (optional) 51 | 52 | - Format: `user_id \t u__ \t ...` 53 | - Optional, only needed for context-aware models, where an argument called `include_user_features` is used to control whether to use the user metadata or not. 54 | - `u__` is the attribute of a user, such as gender, age, and so on. The header should start with `u_` and the is set to `c` for categorical features and `f` for dense (float) features. 55 | 56 | **situation metadata** (optional) 57 | - Situation features are appended to each line of interaction in **train.csv & test.csv & dev.csv** 58 | - Format: `user_id \t item_id \t time \t ... \t c__ \t ...` 59 | - Optional, only needed for context-aware models, where an argument called `include_situation_features` is used to control whether to use the sitaution metadata or not. 60 | - `c__` is the attribute of a situation, such as day of week, hour of day, activity type, and so on. The header should start with `c_` and the is set to `c` for categorical features and `f` for dense (float) features. 61 | 62 | ↓ Examples of different data formats 63 | 64 | ![data format](../docs/_static/data_format.png) 65 | -------------------------------------------------------------------------------- /docs/Getting_Started.md: -------------------------------------------------------------------------------- 1 | ## Getting Started 2 | 3 | 1. Install [Anaconda](https://docs.conda.io/en/latest/miniconda.html) with Python >= 3.10 4 | 2. Clone the repository 5 | 6 | ```bash 7 | git clone https://github.com/THUwangcy/ReChorus.git 8 | ``` 9 | 10 | 3. Install requirements and step into the `src` folder 11 | 12 | ```bash 13 | cd ReChorus 14 | pip install -r requirements.txt 15 | cd src 16 | ``` 17 | 18 | 4. Run model with the build-in dataset 19 | 20 | ```bash 21 | python main.py --model_name BPRMF --emb_size 64 --lr 1e-3 --l2 1e-6 --dataset Grocery_and_Gourmet_Food 22 | ``` 23 | 24 | 5. (optional) Run jupyter notebook in `dataset` folder to download and build new datasets, or prepare your own datasets according to [Guideline](https://github.com/THUwangcy/ReChorus/tree/master/data/README.md) in `data` 25 | 26 | 6. (optional) Implement your own models according to [Guideline](https://github.com/THUwangcy/ReChorus/tree/master/src/README.md) in `src` 27 | -------------------------------------------------------------------------------- /docs/Main_Arguments.md: -------------------------------------------------------------------------------- 1 | ## Main Arguments 2 | 3 | The main arguments are listed below. 4 | 5 | | Args | Default | Description | 6 | | --------------- | --------- | -------------------------- | 7 | | model_name | 'BPRMF' | The name of the model class. | 8 | | model_mode | '' | The task mode for the model to implement. | 9 | | lr | 1e-3 | Learning rate. | 10 | | l2 | 0 | Weight decay in optimizer. | 11 | | test_all | 0 | Wheter to rank all the items during evaluation. (only work in Top-K recommendation tasks)| 12 | | metrics | 'NDCG,HR' | The list of evaluation metrics (seperated by comma). | 13 | | topk | '5,10,20' | The list of K in evaluation metrics (seperated by comma). | 14 | | num_workers | 5 | Number of processes when preparing batches. | 15 | | batch_size | 256 | Batch size during training. | 16 | | eval_batch_size | 256 | Batch size during inference. | 17 | | load | 0 | Whether to load model checkpoint and continue to train. | 18 | | train | 1 | Wheter to perform model training. | 19 | | regenerate | 0 | Wheter to regenerate intermediate files. | 20 | | random_seed | 0 | Random seed of everything. | 21 | | gpu | '0' | The visible GPU device (pass an empty string '' to only use CPU). | 22 | | buffer | 1 | Whether to buffer batches for dev/test. | 23 | | history_max | 20 | The maximum length of history for sequential models. | 24 | | num_neg | 1 | The number of negative items for each training instance. | 25 | | test_epoch | -1 | Print test set metrics every test_epoch during training (-1: no print). | -------------------------------------------------------------------------------- /docs/Supported_Models.md: -------------------------------------------------------------------------------- 1 | ## Supported Models 2 | 3 | We have implemented the following methods (still updating): 4 | 5 | **General Recommenders** 6 | 7 | - [Bayesian personalized ranking from implicit feedback](https://arxiv.org/pdf/1205.2618.pdf?source=post_page) (BPRMF [UAI'09]) 8 | - [Neural Collaborative Filtering](https://arxiv.org/pdf/1708.05031.pdf?source=post_page---------------------------) (NeuMF [WWW'17]) 9 | - [Learning over Knowledge-Base Embeddings for Recommendation](https://arxiv.org/pdf/1803.06540.pdf) (CFKG [SIGIR'18]) 10 | - [LightGCN: Simplifying and Powering Graph Convolution Network for Recommendation](https://dl.acm.org/doi/abs/10.1145/3397271.3401063?casa_token=mMzWDMq9WxQAAAAA%3AsUQEeXtBSLjctZa7qfyOO25nOBqdHWW8ukbjZUeOmcprZcmF3QBWKBtdICrMDidOy8MJ28n3Z1zy5g) (LightGCN [SIGIR'20]) 11 | - [Bootstrapping User and Item Representations for One-Class Collaborative Filtering](https://arxiv.org/pdf/2105.06323) (BUIR [SIGIR'21]) 12 | - [Towards Representation Alignment and Uniformity in Collaborative Filtering](https://arxiv.org/pdf/2206.12811.pdf) (DirectAU [KDD'22]) 13 | 14 | **Sequential Recommenders** 15 | 16 | - [Factorizing Personalized Markov Chains for Next-Basket Recommendation](https://dl.acm.org/doi/pdf/10.1145/1772690.1772773?casa_token=hhM2wEArOQEAAAAA:r_vhs7X8VE0rJ7FF5aZ4i-P-z1mSlBABdw5O9p0cuOahTOQ8D3FVyX6_d58sbQFiV1q1vdVHB-wKqw) (FPMC [WWW'10]) 17 | - [Session-based Recommendations with Recurrent Neural Networks](https://arxiv.org/pdf/1511.06939) (GRU4Rec [ICLR'16]) 18 | - [Neural Attentive Session-based Recommendation](https://arxiv.org/pdf/1711.04725.pdf) (NARM [CIKM'17]) 19 | - [Personalized Top-N Sequential Recommendation via Convolutional Sequence Embedding](https://arxiv.org/pdf/1809.07426) (Caser [WSDM'18]) 20 | - [Self-attentive Sequential Recommendation](https://arxiv.org/pdf/1808.09781.pdf) (SASRec [IEEE'18]) 21 | - [Modeling Item-specific Temporal Dynamics of Repeat Consumption for Recommender Systems](https://dl.acm.org/doi/pdf/10.1145/3308558.3313594) (SLRC [WWW'19]) 22 | - [Time Interval Aware Self-Attention for Sequential Recommendation](https://dl.acm.org/doi/pdf/10.1145/3336191.3371786) (TiSASRec [WSDM'20]) 23 | - [Make It a Chorus: Knowledge- and Time-aware Item Modeling for Sequential Recommendation](http://www.thuir.cn/group/~mzhang/publications/SIGIR2020Wangcy.pdf) (Chorus [SIGIR'20]) 24 | - [Controllable Multi-Interest Framework for Recommendation](https://dl.acm.org/doi/pdf/10.1145/3394486.3403344?casa_token=r35exDCLzSsAAAAA:hbdvRtwvH7LlbllHH7gITV_mpA5hYnAFXcpT2bW8MnbK7Gta50E60xNhC6KoQtY6AGOHaEVsK_GRVQ) (ComiRec [KDD'20]) 25 | - [Towards Dynamic User Intention: Temporal Evolutionary Effects of Item Relations in Sequential Recommendation](https://chenchongthu.github.io/files/TOIS-KDA-wcy.pdf) (KDA [TOIS'21]) 26 | - [Sequential Recommendation with Multiple Contrast Signals](https://dl.acm.org/doi/pdf/10.1145/3522673) (ContraRec [TOIS'22]) 27 | - [Target Interest Distillation for Multi-Interest Recommendation]() (TiMiRec [CIKM'22]) 28 | 29 | **Context-aware Recommenders** 30 | 31 | General Context-aware Recommenders 32 | - [Factorization Machines](https://ieeexplore.ieee.org/document/5694074) (FM [ICDM'10]) 33 | - [Wide {\&} Deep Learning for Recommender Systems](https://dl.acm.org/doi/pdf/10.1145/2988450.2988454) (WideDeep [DLRS@Recsys'16]) 34 | - [Attentional Factorization Machines: Learning the Weight of Feature Interactions via Attention Networks](https://arxiv.org/pdf/1708.04617) (AFM [IJCAI'17]) 35 | - [DeepFM: A Factorization-Machine based Neural Network for CTR Prediction](https://arxiv.org/pdf/1703.04247) (DeepFM [IJCAI'17]) 36 | - [Deep & Cross Network for Ad Click Predictions](https://dl.acm.org/doi/pdf/10.1145/3124749.3124754) (DCN [KDD'17]) 37 | - [AutoInt: Automatic Feature Interaction Learning via Self-Attentive Neural Networks](https://arxiv.org/pdf/1810.11921) (AutoInt [CIKM'18]) 38 | - [xdeepfm: Combining explicit and implicit feature interactions for recommender systems](https://arxiv.org/pdf/1803.05170.pdf%E2%80%8B%E2%80%8B) (xDeepFM [KDD'18]) 39 | - [DCN v2: Improved deep & cross network and practical lessons for web-scale learning to rank systems](https://arxiv.org/pdf/2008.13535) (DCNv2 [WWW'21]) 40 | - [Looking at CTR Prediction Again: Is Attention All You Need?](https://arxiv.org/pdf/2105.05563) (SAM [SIGIR'21]) 41 | - [FinalMLP: an enhanced two-stream MLP model for CTR prediction](https://ojs.aaai.org/index.php/AAAI/article/download/25577/25349) (FinalMLP [AAAI'23]) 42 | 43 | Sequential Context-aware Recommenders 44 | - [Deep interest network for click-through rate prediction](https://arxiv.org/pdf/1706.06978) (DIN [KDD'18]) 45 | - [End-to-end user behavior retrieval in click-through rateprediction model](https://arxiv.org/pdf/2108.04468) (ETA [CoRR'18]) 46 | - [Deep interest evolution network for click-through rate prediction](https://aaai.org/ojs/index.php/AAAI/article/view/4545/4423) (DIEN [AAAI'19]) 47 | - [CAN: feature co-action network for click-through rate prediction](https://dl.acm.org/doi/abs/10.1145/3488560.3498435) (CAN [WSDM'22]) 48 | - [Sampling is all you need on modeling long-term user behaviors for CTR prediction](https://arxiv.org/pdf/2205.10249) (SDIM[CIKM'22]) 49 | 50 | **Impression-based Re-ranking Models** 51 | - [Personalized Re-ranking for Recommendatio](https://arxiv.org/pdf/1904.06813) (PRM [RecSys'19]) 52 | - [SetRank: Learning a Permutation-Invariant Ranking Model for Information Retrieval](https://arxiv.org/pdf/1912.05891) (SIGIR [SIGIR'20]) 53 | - [Multi-Level Interaction Reranking with User Behavior History](https://arxiv.org/pdf/2204.09370) (MIR[SIGIR'22]) -------------------------------------------------------------------------------- /docs/_static/ctr.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/THUwangcy/ReChorus/e58d752558634d42be60de561d6f74df748583cc/docs/_static/ctr.png -------------------------------------------------------------------------------- /docs/_static/data_format.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/THUwangcy/ReChorus/e58d752558634d42be60de561d6f74df748583cc/docs/_static/data_format.png -------------------------------------------------------------------------------- /docs/_static/format_meta.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/THUwangcy/ReChorus/e58d752558634d42be60de561d6f74df748583cc/docs/_static/format_meta.png -------------------------------------------------------------------------------- /docs/_static/format_test.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/THUwangcy/ReChorus/e58d752558634d42be60de561d6f74df748583cc/docs/_static/format_test.png -------------------------------------------------------------------------------- /docs/_static/logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/THUwangcy/ReChorus/e58d752558634d42be60de561d6f74df748583cc/docs/_static/logo.png -------------------------------------------------------------------------------- /docs/_static/logo2.0.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/THUwangcy/ReChorus/e58d752558634d42be60de561d6f74df748583cc/docs/_static/logo2.0.png -------------------------------------------------------------------------------- /docs/_static/module.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/THUwangcy/ReChorus/e58d752558634d42be60de561d6f74df748583cc/docs/_static/module.png -------------------------------------------------------------------------------- /docs/_static/module_new.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/THUwangcy/ReChorus/e58d752558634d42be60de561d6f74df748583cc/docs/_static/module_new.png -------------------------------------------------------------------------------- /docs/_static/rerank.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/THUwangcy/ReChorus/e58d752558634d42be60de561d6f74df748583cc/docs/_static/rerank.png -------------------------------------------------------------------------------- /docs/_static/topk.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/THUwangcy/ReChorus/e58d752558634d42be60de561d6f74df748583cc/docs/_static/topk.png -------------------------------------------------------------------------------- /docs/demo_scripts_results/CTR_MIND.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh -x 2 | 3 | # CTR prediction on MIND dataset 4 | python main.py --model_name FM --lr 5e-4 --l2 0 --dataset MINDCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 5 | 6 | python main.py --model_name WideDeep --lr 5e-4 --l2 0 --dropout 0.5 --layers "[64,64,64]" --dataset MINDCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 7 | 8 | python main.py --model_name DeepFM --lr 5e-4 --l2 0 --dropout 0.2 --layers "[512,64]" --dataset MINDCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 9 | 10 | python main.py --model_name AFM --lr 5e-4 --l2 0 --dropout 0 --attention_size 64 --reg_weight 1.0 --dataset MINDCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 11 | 12 | python main.py --model_name DCN --lr 5e-4 --l2 0 --layers "[128,64]" --cross_layer_num 4 --reg_weight 1.0 --dropout 0.8 --dataset MINDCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 13 | 14 | python main.py --model_name xDeepFM --lr 1e-4 --l2 0 --layers "[128,64]" --cin_layers "[8,8]" --direct 0 --reg_weight 0 --dataset MINDCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 15 | 16 | python main.py --model_name AutoInt --lr 5e-4 --l2 0 --dropout 0.5 --attention_size 64 --num_heads 2 --num_layers 2 --layers "[64,64,64]" --dataset MINDCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 17 | 18 | python main.py --model_name DCNv2 --lr 2e-3 --l2 0 --layers "[256,256,256]" --cross_layer_num 3 --mixed 1 --structure stacked --low_rank 64 --expert_num 2 --reg_weight 2.0 --dataset MINDCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 19 | 20 | python main.py --model_name FinalMLP --mlp1_dropout 0 --mlp2_dropout 0 --mlp1_batch_norm 0 --mlp2_batch_norm 0 --use_fs 1 --lr 5e-3 --l2 0 --dataset MINDCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE --fs1_context c_hour_c,c_weekday_c,c_period_c,c_day_f --fs2_context i_category_c,i_subcategory_c --mlp1_hidden_units "[64]" --mlp2_hidden_units "[256]" --fs_hidden_units "[64]" 21 | 22 | python main.py --model_name SAM --lr 5e-4 --l2 0 --interaction_type SAM3A --aggregation mean_pooling --num_layers 5 --use_residual 1 --dropout 0 --dataset MINDCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 23 | 24 | python main.py --model_name DIN --dropout 0 --lr 5e-3 --l2 0 --history_max 40 --dataset MINDCTR --path ../data/ --num_neg 0 --batch_size 1024 --eval_batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --att_layers "[64,64]" --dnn_layers "[64]" --loss_n BCE 25 | 26 | python main.py --model_name DIEN --lr 5e-4 --l2 0 --history_max 30 --alpha_aux 0.1 --aux_hidden_layers "[64]" --fcn_hidden_layers "[64]" --evolving_gru_type AIGRU --dropout 0 --dataset MINDCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 27 | 28 | python main.py --model_name CAN --lr 5e-4 --l2 0 --co_action_layers "[4,4]" --orders 1 --induce_vec_size 512 --history_max 20 --alpha_aux 0.1 --aux_hidden_layers "[64]" --fcn_hidden_layers "[64]" --evolving_gru_type AGRU --dropout 0 --dataset MINDCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE -------------------------------------------------------------------------------- /docs/demo_scripts_results/CTR_ML1M.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh -x 2 | 3 | # CTR prediction on MovieLens-1M dataset 4 | python main.py --model_name FM --lr 1e-3 --l2 1e-4 --dataset ML_1MCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 5 | 6 | python main.py --model_name WideDeep --lr 5e-3 --l2 0 --dropout 0.5 --layers "[64,64,64]" --dataset ML_1MCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 7 | 8 | python main.py --model_name DeepFM --lr 1e-3 --l2 1e-4 --dropout 0.2 --layers "[512,128]" --dataset ML_1MCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 9 | 10 | python main.py --model_name AFM --lr 5e-4 --l2 1e-4 --dropout 0.8 --attention_size 128 --reg_weight 0.5 --dataset ML_1MCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 11 | 12 | python main.py --model_name DCN --lr 5e-4 --l2 1e-4 --layers "[512,128]" --cross_layer_num 1 --reg_weight 0.5 --dataset ML_1MCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 13 | 14 | python main.py --model_name xDeepFM --lr 1e-3 --l2 1e-4 --layers "[512,512,512]" --cin_layers "[8,8]" --direct 0 --reg_weight 0 --dataset ML_1MCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 15 | 16 | python main.py --model_name AutoInt --lr 2e-3 --l2 1e-6 --dropout 0.2 --attention_size 64 --num_heads 2 --num_layers 2 --layers "[64,64,64]" --dataset ML_1MCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 17 | 18 | python main.py --model_name DCNv2 --lr 1e-3 --l2 1e-4 --layers "[256,256,256]" --cross_layer_num 3 --mixed 0 --structure parallel --low_rank 64 --expert_num 1 --reg_weight 2.0 --dataset ML_1MCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 19 | 20 | python main.py --model_name FinalMLP --mlp1_dropout 0.2 --mlp2_dropout 0.5 --mlp1_batch_norm 1 --mlp2_batch_norm 1 --use_fs 1 --lr 5e-3 --l2 1e-6 --dataset ML_1MCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE --fs1_context c_hour_c,c_weekday_c,c_period_c,c_day_f --fs2_context i_genre_c,i_title_c --mlp1_hidden_units "[64]" --mlp2_hidden_units "[64,64]" --fs_hidden_units "[256,64]" 21 | 22 | python main.py --model_name SAM --lr 1e-3 --l2 1e-4 --interaction_type SAM3A --aggregation mean_pooling --num_layers 1 --use_residual 0 --dropout 0.5 --dataset ML_1MCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 23 | 24 | python main.py --model_name DIN --history_max 20 --lr 5e-4 --l2 1e-4 --dnn_layers "[512,64]" --att_layers "[64]" --dropout 0.5 --dataset ML_1MCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 25 | 26 | python main.py --model_name DIEN --lr 5e-3 --l2 1e-6 --history_max 20 --alpha_aux 0.5 --aux_hidden_layers "[64,64,64]" --fcn_hidden_layers "[256]" --evolving_gru_type AIGRU --dropout 0.2 --dataset ML_1MCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 27 | 28 | python main.py --model_name CAN --lr 2e-3 --l2 1e-4 --co_action_layers "[4,4,4]" --orders 1 --induce_vec_size 1024 --history_max 30 --alpha_aux 0.1 --aux_hidden_layers "[64,64,64]" --fcn_hidden_layers "[256,128]" --evolving_gru_type AIGRU --dropout 0.2 --dataset ML_1MCTR --path ../data/ --num_neg 0 --batch_size 1024 --metric AUC,Log_loss --include_item_features 1 --include_situation_features 1 --model_mode CTR --loss_n BCE 29 | 30 | -------------------------------------------------------------------------------- /docs/demo_scripts_results/README.md: -------------------------------------------------------------------------------- 1 | # Demo Scripts and Results 2 | 3 | The demo scripts and corresponding results are provided in this directory. 4 | 5 | ## Impression-based Ranking and Reranking 6 | 7 | The results of general, sequential, and reranking models in `MovieLens-1M` and `MIND-Large` datasets. 8 | Training, validation, and test sets are split along the global timeline: In MIND, the first six days are treated as training set, followed by half-day validation set and half-day test set; In MovieLens-1M, training, validation, and test sets are split with 80\%, 10\%, 10\% of the time. 9 | The detailed preprocessing scripts are shown in [ML-Preprocessing](https://github.com/THUwangcy/ReChorus/tree/master/data/MovieLens_1M/MovieLens-1M.ipynb) and [MIND-Preprocessing](https://github.com/THUwangcy/ReChorus/tree/master/data/MIND_Large/MIND-large.ipynb). 10 | 11 | The configurations are shown in `Rerank_ML1M.sh` and `Rerank_MIND.sh`. 12 | 13 | ![rerankResults](../_static/rerank.png) 14 | 15 | ## Top-k Recommendation with Context-aware models 16 | 17 | The results of context-aware models in `MovieLens-1M` and `MIND-Large` datasets for the Top-k recommendation task. 18 | Training, validation, and test sets are split along the global timeline: In MIND, the first six days are treated as training set, followed by half-day validation set and half-day test set; In MovieLens-1M, training, validation, and test sets are split with 80\%, 10\%, 10\% of the time. 19 | We randomly sample 99 negative items for each test case to rank together with the ground-truth item (also support ranking over all the items with `--test_all 1`). 20 | 21 | The configurations are shown in `Topk_ML1M.sh` and `Topk_MIND.sh`. 22 | 23 | 24 | ![topkResults](../_static/topk.png) 25 | 26 | ## CTR Prediction with Context-aware models 27 | 28 | The results of context-aware models in `MovieLens-1M` and `MIND-Large` datasets for the CTR prediction task. 29 | Training, validation, and test sets are split along the global timeline: In MIND, the first six days are treated as training set, followed by half-day validation set and half-day test set; In MovieLens-1M, training, validation, and test sets are split with 80\%, 10\%, 10\% of the time. 30 | 31 | The configurations are shown in `CTR_ML1M.sh` and `CTR_MIND.sh`. 32 | 33 | ![ctrResults](../_static/ctr.png) 34 | 35 | ## Top-k Recommendation on Amazon dataset 36 | > From ReChorus1.0 37 | 38 | The results of general and sequential models in `Grocery_and_Gourmet_Food` dataset (151.3k entries). 39 | Leave-one-out is applied to split data: the most recent interaction of each user for testing, the second recent item for validation, and the remaining items for training. 40 | We randomly sample 99 negative items for each test case to rank together with the ground-truth item (also support ranking over all the items with `--test_all 1`). 41 | 42 | The configurations are shown in `Topk_Amazon.sh` 43 | 44 | | Model | HR@5 | NDCG@5 | Time/iter | Sequential | Knowledge | Time-aware | 45 | |:------------------------------------------------------------------------------------------------- |:------:|:------:|:---------:|:----------:|:---------:|:----------:| 46 | | [MostPop](https://github.com/THUwangcy/ReChorus/tree/master/src/models/general/POP.py) | 0.2065 | 0.1301 | - | | | | 47 | | [BPRMF](https://github.com/THUwangcy/ReChorus/tree/master/src/models/general/BPRMF.py) | 0.3549 | 0.2486 | 2.5s | | | | 48 | | [NeuMF](https://github.com/THUwangcy/ReChorus/tree/master/src/models/general/NCF.py) | 0.3237 | 0.2221 | 3.4s | | | | 49 | | [LightGCN](https://github.com/THUwangcy/ReChorus/tree/master/src/models/general/LightGCN.py) | 0.3705 | 0.2564 | 6.1s | | | | 50 | | [BUIR](https://github.com/THUwangcy/ReChorus/tree/master/src/models/general/BUIR.py) | 0.3701 | 0.2567 | 3.3s | | | | 51 | | [DirectAU](https://github.com/THUwangcy/ReChorus/tree/master/src/models/general/DirectAU.py) | 0.3911 | 0.2779 | 3.3s | | | | 52 | | [FPMC](https://github.com/THUwangcy/ReChorus/tree/master/src/models/sequential/FPMC.py) | 0.3594 | 0.2785 | 3.4s | √ | | | 53 | | [GRU4Rec](https://github.com/THUwangcy/ReChorus/tree/master/src/models/sequential/GRU4Rec.py) | 0.3659 | 0.2614 | 4.9s | √ | | | 54 | | [NARM](https://github.com/THUwangcy/ReChorus/tree/master/src/models/sequential/NARM.py) | 0.3650 | 0.2617 | 7.5s | √ | | | 55 | | [Caser](https://github.com/THUwangcy/ReChorus/tree/master/src/models/sequential/Caser.py) | 0.3526 | 0.2499 | 7.8s | √ | | | 56 | | [SASRec](https://github.com/THUwangcy/ReChorus/tree/master/src/models/sequential/SASRec.py) | 0.3917 | 0.2942 | 5.5s | √ | | | 57 | | [ComiRec](https://github.com/THUwangcy/ReChorus/tree/master/src/models/sequential/ComiRec.py) | 0.3753 | 0.2675 | 4.5s | √ | | | 58 | | [TiMiRec+](https://github.com/THUwangcy/ReChorus/tree/master/src/models/sequential/TiMiRec.py) | 0.4020 | 0.3016 | 8.8s | √ | | | 59 | | [ContraRec](https://github.com/THUwangcy/ReChorus/tree/master/src/models/sequential/ContraRec.py) | 0.4251 | 0.3285 | 5.6s | √ | | | 60 | | [TiSASRec](https://github.com/THUwangcy/ReChorus/tree/master/src/models/sequential/TiSASRec.py) | 0.3949 | 0.2945 | 7.6s | √ | | √ | 61 | | [CFKG](https://github.com/THUwangcy/ReChorus/tree/master/src/models/general/CFKG.py) | 0.4199 | 0.2984 | 8.7s | | √ | | 62 | | [SLRC+](https://github.com/THUwangcy/ReChorus/tree/master/src/models/sequential/SLRCPlus.py) | 0.4376 | 0.3263 | 4.3s | √ | √ | √ | 63 | | [Chorus](https://github.com/THUwangcy/ReChorus/tree/master/src/models/sequential/Chorus.py) | 0.4668 | 0.3414 | 4.9s | √ | √ | √ | 64 | | [KDA](https://github.com/THUwangcy/ReChorus/tree/master/src/models/sequential/KDA.py) | 0.5191 | 0.3901 | 9.9s | √ | √ | √ | 65 | | [ContraKDA](https://github.com/THUwangcy/ReChorus/tree/master/src/models/sequential/ContraKDA.py) | 0.5282 | 0.3992 | 13.6s | √ | √ | √ | 66 | 67 | -------------------------------------------------------------------------------- /docs/demo_scripts_results/Rerank_MIND.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh -x 2 | 3 | # Impression-based Ranking and Reranking on MIND dataset 4 | python main.py --model_name BPRMF --emb_size 64 --lr 1e-3 --l2 0 --loss_n BPR --dataset MINDCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10 --main_metric NDCG@2 --model_mode Impression 5 | 6 | python main.py --model_name GRU4Rec --hidden_size 64 --history_max 20 --emb_size 64 --lr 2e-3 --l2 0 --loss_n BPR --dataset MINDCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10 --main_metric NDCG@2 --model_mode Impression 7 | 8 | python main.py --model_name SASRec --num_layers 3 --num_heads 1 --history_max 20 --emb_size 64 --lr 5e-4 --l2 0 --loss_n BPR --dataset MINDCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10 --main_metric NDCG@2 --model_mode Impression 9 | 10 | # reranking model: the corresponding .yaml config file and .pt backbone checkpoint should be prepared in ../model/{ranker_name}/ in advance 11 | random_seed=0 12 | # PRM+BPRMF 13 | python main.py --model_name PRM --num_hidden_unit 64 --positionafter 0 --emb_size 64 --n_blocks 4 --num_heads 4 --lr 5e-4 --l2 0 --loss_n BPR --dataset MINDCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10,20 --main_metric NDCG@2 --num_workers 0 --ranker_name BPRMF --ranker_config_file BPRMF_best.yaml --ranker_model_file BPRMF__MINDCTR__${random_seed}__best.pt --model_mode General 14 | 15 | # SetRank+BPRMF 16 | python main.py --model_name SetRank --emb_size 64 --n_blocks 4 --num_heads 4 --num_hidden_unit 64 --setrank_type IMSAB --positionafter 1 --lr 1e-3 --l2 0 --loss_n BPR --dataset MINDCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10 --main_metric NDCG@2 --num_workers 0 --ranker_name BPRMF --ranker_config_file BPRMF_best.yaml --ranker_model_file BPRMF__MINDCTR__${random_seed}__best.pt --model_mode General 17 | 18 | # MIR+BPRMF 19 | python main.py --model_name MIR --emb_size 64 --num_heads 4 --num_hidden_unit 64 --lr 5e-4 --l2 0 --loss_n BPR --dataset MINDCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10 --main_metric NDCG@2 --num_workers 0 --ranker_name BPRMF --ranker_config_file BPRMF_best.yaml --ranker_model_file BPRMF__MINDCTR__${random_seed}__best.pt --model_mode General 20 | 21 | # PRM+GRU4Rec 22 | python main.py --model_name PRM --num_hidden_unit 64 --history_max 10 --positionafter 1 --emb_size 64 --n_blocks 6 --num_heads 4 --lr 2e-3 --l2 0 --loss_n BPR --dataset MINDCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10,20 --main_metric NDCG@2 --num_workers 0 --ranker_name GRU4Rec --ranker_config_file GRU4Rec_MINDCTR_best.yaml --ranker_model_file GRU4Rec__MINDCTR__${random_seed}__best.pt --model_mode Sequential 23 | 24 | # SetRank+GRU4Rec 25 | python main.py --model_name SetRank --emb_size 64 --history_max 20 --n_blocks 8 --num_heads 4 --num_hidden_unit 32 --setrank_type MSAB --positionafter 1 --lr 1e-3 --l2 0 --loss_n BPR --dataset MINDCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10 --main_metric NDCG@2 --num_workers 0 --ranker_name GRU4Rec --ranker_config_file GRU4Rec_MINDCTR_best.yaml --ranker_model_file GRU4Rec__MINDCTR__${random_seed}__best.pt --model_mode Sequential 26 | 27 | # MIR+GRU4Rec 28 | python main.py --model_name MIR --history_max 10 --emb_size 64 --num_heads 2 --num_hidden_unit 64 --lr 1e-3 --l2 0 --loss_n BPR --dataset MINDCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10 --main_metric NDCG@2 --num_workers 0 --ranker_name GRU4Rec --ranker_config_file GRU4Rec_MINDCTR_best.yaml --ranker_model_file GRU4Rec__MINDCTR__${random_seed}__best.pt --model_mode Sequential -------------------------------------------------------------------------------- /docs/demo_scripts_results/Rerank_ML1M.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh -x 2 | 3 | # Impression-based Ranking and Reranking on MovieLens-1M dataset 4 | python main.py --model_name BPRMF --emb_size 64 --lr 1e-3 --l2 0 --loss_n BPR --dataset ML_1MCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10 --main_metric NDCG@2 --model_mode Impression 5 | 6 | python main.py --model_name LightGCN --n_layers 1 --emb_size 64 --lr 1e-3 --l2 0 --loss_n BPR --dataset ML_1MCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10 --main_metric NDCG@2 --model_mode Impression 7 | 8 | python main.py --model_name GRU4Rec --hidden_size 32 --history_max 30 --emb_size 64 --lr 1e-3 --l2 1e-6 --loss_n BPR --dataset ML_1MCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10 --main_metric NDCG@2 --model_mode Impression 9 | 10 | python main.py --model_name SASRec --num_layers 3 --num_heads 2 --history_max 20 --emb_size 64 --lr 2e-4 --l2 1e-6 --loss_n BPR --dataset ML_1MCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10 --main_metric NDCG@2 --model_mode Impression 11 | 12 | # reranking model: the corresponding .yaml config file and .pt backbone checkpoint should be prepared in ../model/{ranker_name}/ in advance 13 | random_seed=0 14 | # PRM+LightGCN 15 | python main.py --model_name PRM --positionafter 1 --num_hidden_unit 256 --emb_size 64 --n_blocks 4 --num_heads 2 --lr 1e-3 --l2 1e-6 --loss_n BPR --dataset ML_1MCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10,20 --main_metric NDCG@2 --num_workers 0 --ranker_name LightGCN --ranker_config_file ml_LightGCN_best.yaml --ranker_model_file LightGCN__ML_1MCTR__${random_seed}__best.pt --model_mode General 16 | 17 | # SetRank+LightGCN 18 | python main.py --model_name SetRank --emb_size 64 --n_blocks 4 --num_heads 4 --num_hidden_unit 32 --setrank_type IMSAB --positionafter 1 --lr 2e-4 --l2 1e-4 --loss_n BPR --dataset ML_1MCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10 --main_metric NDCG@2 --num_workers 0 --ranker_name LightGCN --ranker_config_file ml_LightGCN_best.yaml --ranker_model_file LightGCN__ML_1MCTR__${random_seed}__best.pt --model_mode General 19 | 20 | # MIR+LightGCN 21 | python main.py --model_name MIR --emb_size 64 --history_max 10 --num_heads 2 --num_hidden_unit 32 --lr 2e-3 --l2 0 --loss_n BPR --dataset ML_1MCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10 --main_metric NDCG@2 --num_workers 0 --ranker_name LightGCN --ranker_config_file ml_LightGCN_best.yaml --ranker_model_file LightGCN__ML_1MCTR__${random_seed}__best.pt --model_mode General 22 | 23 | # PRM+SASRec 24 | python main.py --model_name PRM --emb_size 64 --history_max 20 --n_blocks 1 --num_heads 1 --lr 1e-3 --l2 0 --loss_n BPR --dataset ML_1MCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10,20 --main_metric NDCG@2 --num_workers 0 --ranker_name SASRec --ranker_config_file ml_SASRec_best.yaml --ranker_model_file SASRec__ML_1MCTR__${random_seed}__best.pt --model_mode Sequential 25 | 26 | # SetRank+SASRec 27 | python main.py --model_name SetRank --emb_size 64 --history_max 20 --n_blocks 4 --num_heads 1 --num_hidden_unit 64 --setrank_type MSAB --positionafter 1 --lr 2e-4 --l2 1e-6 --loss_n BPR --dataset ML_1MCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10 --main_metric NDCG@2 --num_workers 0 --ranker_name SASRec_test --ranker_config_file ml_SASRec_best.yaml --ranker_model_file SASRec_test__ML_1MCTR__${random_seed}__best.pt --model_mode Sequential 28 | 29 | # MIR+SASRec 30 | python main.py --model_name MIR --history_max 20 --emb_size 64 --num_heads 4 --num_hidden_unit 128 --lr 2e-4 --l2 1e-6 --loss_n BPR --dataset ML_1MCTR --path ../data/ --metric NDCG,HR --topk 1,2,3,5,10 --main_metric NDCG@2 --num_workers 0 --ranker_name SASRec_test --ranker_config_file ml_SASRec_best.yaml --ranker_model_file SASRec_test__ML_1MCTR__${random_seed}__best.pt --model_mode Sequential 31 | -------------------------------------------------------------------------------- /docs/demo_scripts_results/Topk_Amazon.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh -x 2 | 3 | # Grocery_and_Gourmet_Food 4 | python main.py --model_name POP --train 0 --dataset 'Grocery_and_Gourmet_Food' 5 | 6 | python main.py --model_name BPRMF --emb_size 64 --lr 1e-3 --l2 1e-6 --dataset 'Grocery_and_Gourmet_Food' 7 | 8 | python main.py --model_name NeuMF --emb_size 64 --layers '[64]' --lr 5e-4 --l2 1e-7 --dropout 0.2 --dataset 'Grocery_and_Gourmet_Food' 9 | 10 | python main.py --model_name CFKG --emb_size 64 --margin 1 --include_attr 1 --lr 1e-4 --l2 1e-8 --dataset 'Grocery_and_Gourmet_Food' 11 | 12 | python main.py --model_name LightGCN --emb_size 64 --n_layers 3 --lr 1e-3 --l2 1e-8 --dataset 'Grocery_and_Gourmet_Food' 13 | 14 | python main.py --model_name BUIR --emb_size 64 --lr 1e-3 --l2 1e-6 --dataset 'Grocery_and_Gourmet_Food' 15 | 16 | python main.py --model_name DirectAU --emb_size 64 --lr 1e-3 --l2 1e-5 --gamma 0.3 --epoch 500 --dataset 'Grocery_and_Gourmet_Food' 17 | 18 | python main.py --model_name FPMC --emb_size 64 --lr 1e-3 --l2 1e-6 --history_max 20 --dataset 'Grocery_and_Gourmet_Food' 19 | 20 | python main.py --model_name GRU4Rec --emb_size 64 --hidden_size 100 --lr 1e-3 --l2 1e-4 --history_max 20 --dataset 'Grocery_and_Gourmet_Food' 21 | 22 | python main.py --model_name NARM --emb_size 64 --hidden_size 100 --attention_size 4 --lr 1e-3 --l2 1e-4 --history_max 20 --dataset 'Grocery_and_Gourmet_Food' 23 | 24 | python main.py --model_name Caser --emb_size 64 --L 5 --num_horizon 64 --num_vertical 32 --lr 1e-3 --l2 1e-4 --history_max 20 --dataset 'Grocery_and_Gourmet_Food' 25 | 26 | python main.py --model_name SASRec --emb_size 64 --num_layers 1 --num_heads 1 --lr 1e-4 --l2 1e-6 --history_max 20 --dataset 'Grocery_and_Gourmet_Food' 27 | 28 | python main.py --model_name SLRCPlus --emb_size 64 --lr 5e-4 --l2 1e-5 --dataset 'Grocery_and_Gourmet_Food' 29 | 30 | python main.py --model_name TiSASRec --emb_size 64 --num_layers 1 --num_heads 1 --lr 1e-4 --l2 1e-6 --history_max 20 --dataset 'Grocery_and_Gourmet_Food' 31 | 32 | python main.py --model_name Chorus --emb_size 64 --margin 1 --lr 5e-4 --l2 1e-5 --epoch 50 --early_stop 0 --batch_size 512 --dataset 'Grocery_and_Gourmet_Food' --stage 1 33 | python main.py --model_name Chorus --emb_size 64 --margin 1 --lr_scale 0.1 --lr 1e-3 --l2 0 --dataset 'Grocery_and_Gourmet_Food' --base_method 'BPR' --stage 2 34 | 35 | python main.py --model_name ComiRec --emb_size 64 --lr 1e-3 --l2 1e-6 --attn_size 8 --K 4 --add_pos 1 --history_max 20 --dataset 'Grocery_and_Gourmet_Food' 36 | 37 | python main.py --model_name KDA --emb_size 64 --include_attr 1 --freq_rand 0 --lr 1e-3 --l2 1e-6 --num_heads 4 --history_max 20 --dataset 'Grocery_and_Gourmet_Food' 38 | 39 | python main.py --model_name ContraRec --emb_size 64 --lr 1e-4 --l2 1e-6 --history_max 20 --encoder BERT4Rec --gamma 1 --temp 0.2 --batch_size 4096 --dataset 'Grocery_and_Gourmet_Food' 40 | 41 | python main.py --model_name TiMiRec --emb_size 64 --lr 1e-4 --l2 1e-6 --history_max 20 --K 6 --add_pos 1 --add_trm 1 --stage pretrain --dataset 'Grocery_and_Gourmet_Food' 42 | python main.py --model_name TiMiRec --emb_size 64 --lr 1e-4 --l2 1e-6 --history_max 20 --K 6 --add_pos 1 --add_trm 1 --stage finetune --temp 1 --n_layers 1 --check_epoch 10 --dataset 'Grocery_and_Gourmet_Food' -------------------------------------------------------------------------------- /docs/demo_scripts_results/Topk_MIND.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh -x 2 | 3 | # Top-k recommendation on MIND dataset 4 | python main.py --model_name FM --lr 1e-3 --l2 1e-4 --dataset MINDTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 5 | 6 | python main.py --model_name WideDeep --lr 5e-4 --l2 0 --dropout 0.5 --layers "[64,64,64]" --dataset MINDTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 7 | 8 | python main.py --model_name DeepFM --lr 2e-3 --l2 1e-4 --dropout 0.5 --layers "[64]" --dataset MINDTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 9 | 10 | python main.py --model_name AFM --lr 2e-3 --l2 1e-6 --dropout 0.8 --attention_size 64 --reg_weight 0 --dataset MINDTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 11 | 12 | python main.py --model_name DCN --lr 2e-3 --l2 1e-4 --layers "[512,128]" --cross_layer_num 5 --reg_weight 2.0 --dataset MINDTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 13 | 14 | python main.py --model_name xDeepFM --lr 2e-3 --l2 1e-4 --layers "[64,64,64]" --cin_layers "[8,8]" --direct 0 --reg_weight 2.0 --dataset MINDTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 15 | 16 | python main.py --model_name AutoInt --lr 2e-3 --l2 1e-4 --dropout 0.8 --attention_size 64 --num_heads 1 --num_layers 1 --layers "[512]" --dataset MINDTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 17 | 18 | python main.py --model_name DCNv2 --lr 2e-3 --l2 1e-4 --layers "[64,64,64]" --cross_layer_num 4 --mixed 1 --structure parallel --low_rank 64 --expert_num 1 --reg_weight 2.0 --dataset MINDTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 19 | 20 | python main.py --model_name FinalMLP --mlp1_hidden_units "[256,256,256]" --mlp2_hidden_units "[256,128,64]" --mlp1_dropout 0.5 --mlp2_dropout 0.5 --use_fs 0 --mlp1_batch_norm 0 --mlp2_batch_norm 1 --lr 5e-3 --l2 0 --fs1_context c_hour_c,c_weekday_c,c_period_c,c_day_f --fs2_context i_category_c,i_subcategory_c --dataset MINDTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 21 | 22 | python main.py --model_name SAM --lr 2e-4 --l2 1e-4 --interaction_type SAM3A --aggregation mean_pooling --num_layers 1 --use_residual 0 --dropout 0.0 --dataset MINDTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 23 | 24 | python main.py --model_name DIN --lr 2e-3 --l2 0 --history_max 10 --dataset MINDTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK --att_layers "[256]" --dnn_layers "[64,64]" --dropout 0.8 25 | 26 | python main.py --model_name DIEN --lr 5e-4 --l2 1e-4 --history_max 30 --alpha_aux 0.1 --aux_hidden_layers "[64]" --fcn_hidden_layers "[64]" --evolving_gru_type AUGRU --dropout 0 --dataset MINDTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 32 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 27 | 28 | python main.py --model_name CAN --lr 2e-3 --l2 0 --co_action_layers "[4,4]" --orders 1 --induce_vec_size 512 --history_max 30 --alpha_aux 0.1 --aux_hidden_layers "[64]" --fcn_hidden_layers "[64]" --evolving_gru_type AUGRU --dropout 0 --dataset MINDTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 32 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK -------------------------------------------------------------------------------- /docs/demo_scripts_results/Topk_ML1M.sh: -------------------------------------------------------------------------------- 1 | #!/bin/sh -x 2 | 3 | # Top-k recommendation on MovieLens-1M dataset 4 | python main.py --model_name FM --lr 1e-3 --l2 0 --dataset ML_1MTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 5 | 6 | python main.py --model_name WideDeep --lr 1e-3 --l2 0 --dropout 0.5 --layers "[64,64,64]" --dataset ML_1MTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 7 | 8 | python main.py --model_name DeepFM --lr 5e-4 --l2 1e-6 --dropout 0.5 --layers "[512,128]" --dataset ML_1MTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 9 | 10 | python main.py --model_name AFM --lr 5e-3 --l2 0 --dropout 0.5 --attention_size 64 --reg_weight 2.0 --dataset ML_1MTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 11 | 12 | python main.py --model_name DCN --lr 5e-4 --l2 1e-4 --layers "[64,64,64]" --cross_layer_num 2 --reg_weight 0.5 --dataset ML_1MTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 13 | 14 | python main.py --model_name xDeepFM --lr 5e-4 --l2 0 --dropout 0.8 --layers "[512,512,512]" --cin_layers "[8,8]" --direct 0 --reg_weight 1.0 --dataset ML_1MTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 15 | 16 | python main.py --model_name AutoInt --lr 2e-3 --l2 0 --dropout 0 --attention_size 64 --num_heads 2 --num_layers 2 --layers "[256]" --dataset ML_1MTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 17 | 18 | python main.py --model_name DCNv2 --dropout 0 --lr 1e-3 --l2 1e-4 --layers "[256,64]" --cross_layer_num 2 --mixed 0 --structure stacked --low_rank 64 --expert_num 2 --reg_weight 2.0 --dataset ML_1MTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 19 | 20 | python main.py --model_name FinalMLP --mlp1_hidden_units "[64]" --mlp2_hidden_units "[64,64,64]" --mlp1_dropout 0.5 --mlp2_dropout 0.2 --use_fs 1 --mlp1_batch_norm 0 --mlp2_batch_norm 0 --lr 1e-3 --l2 0 --fs1_context c_hour_c,c_weekday_c,c_period_c,c_day_f --fs2_context i_genre_c,i_title_c --dataset ML_1MTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 21 | 22 | python main.py --model_name SAM --lr 1e-3 --l2 1e-4 --interaction_type SAM3A --aggregation mean_pooling --num_layers 1 --use_residual 1 --dropout 0.2 --dataset ML_1MTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 23 | 24 | python main.py --model_name DIN --lr 2e-3 --l2 1e-6 --history_max 10 --dataset ML_1MTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 128 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK --att_layers "[64,64,64]" --dnn_layers "[128,64]" --dropout 0.5 25 | 26 | python main.py --model_name DIEN --lr 5e-4 --l2 1e-6 --history_max 20 --alpha_aux 0.1 --aux_hidden_layers "[64]" --fcn_hidden_layers "[64]" --evolving_gru_type AIGRU --dropout 0 --dataset ML_1MTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 32 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK 27 | 28 | python main.py --model_name CAN --lr 5e-4 --l2 1e-4 --co_action_layers "[4,4]" --orders 2 --induce_vec_size 1024 --history_max 10 --alpha_aux 0.1 --aux_hidden_layers "[64]" --fcn_hidden_layers "[64,64]" --evolving_gru_type AIGRU --dropout 0.2 --dataset ML_1MTOPK --path ../data/ --num_neg 1 --batch_size 256 --eval_batch_size 32 --metric NDCG,HR --topk 3,5,10,20 --include_item_features 1 --include_situation_features 1 --model_mode TopK -------------------------------------------------------------------------------- /docs/tutorials/README.md: -------------------------------------------------------------------------------- 1 | ## Tutorials for 3 Basic Tasks 2 | 3 | * [Impression-based Ranking/Re-ranking Task](https://github.com/THUwangcy/ReChorus/tree/master/docs/tutorials/Impression_based_Ranking_Reranking.ipynb) 4 | * [CTR Prediction Task](https://github.com/THUwangcy/ReChorus/tree/master/docs/tutorials/CTR_Prediction.ipynb) 5 | * [Context-aware Top-k Recommendation Task](https://github.com/THUwangcy/ReChorus/tree/master/docs/tutorials/Context_aware_Topk_Recommendation.ipynb) -------------------------------------------------------------------------------- /requirements.txt: -------------------------------------------------------------------------------- 1 | python==3.10.4 2 | torch==1.12.1 3 | cudatoolkit==10.2.89 4 | numpy==1.22.3 5 | ipython==8.10.0 6 | jupyter==1.0.0 7 | tqdm==4.66.1 8 | pandas==1.4.4 9 | scikit-learn==1.1.3 10 | scipy==1.7.3 11 | pickle 12 | yaml 13 | -------------------------------------------------------------------------------- /src/README.md: -------------------------------------------------------------------------------- 1 | # Source Code 2 | 3 | `main.py` serves as the entrance of our framework, and there are three main packages. 4 | 5 | ### Structure 6 | 7 | - `helpers\` 8 | - `BaseReader.py`: read dataset csv into DataFrame and append necessary information (e.g. interaction history) 9 | - `ContextReader.py`: inherited from BaseReader, read user&item metadata, and count statistics about all context features 10 | - `ContextSeqReader.py`: inherited from ContextReader, append interaction history with situation context features. 11 | - `ImpressionReader.py`: inherited from BaseReader, group interactions with the same impression id into an instance. 12 | - `BaseRunner.py`: control the training and evaluation process of a model 13 | - `CTRRunner.py`: inherited from BaseRunner, train and evaluate a model with binary label. (Click-through-rate Predition task) 14 | - `ImpressionRunner.py`: inherited from BaseRunner, train and evaluate a model with impression-based logs (Variable lengths of positive and negative items in a list). 15 | - `...`: customize helpers with specific functions 16 | - `models\` 17 | - `BaseModel.py`: basic model classes and dataset classes, with some common functions of a model 18 | - `BaseContextModel.py`: inherited from BaseModel, add context features for base model 19 | - `BaseImpressionModel.py`: inherited from BaseModel, construct data batch in impressions 20 | - `...`: customize models inherited from classes in *BaseModel* 21 | - `utils\` 22 | - `layers.py`: common modules for model definition (e.g. attention and MLP blocks) 23 | - `utils.py`: some utils functions 24 | - `main.py`: main entrance, connect all the modules 25 | - `exp.py`: repeat experiments in *run.sh* and save averaged results to csv 26 | 27 | ### Define a New Model 28 | 29 | Generally we can define a new class inheriting *GeneralModel* (a subclass of *BaseModel*), as well as the inner class *Dataset*. The following functions need to be implement at least: 30 | 31 | ```python 32 | class NewModel(GeneralModel): 33 | reader = 'BaseReader' # assign a reader class, BaseReader by default 34 | runner = 'BaseRunner' # assign a runner class, BaseRunner by default 35 | 36 | def __init__(self, args, corpus): 37 | super().__init__(args, corpus) 38 | self._define_params() 39 | self.apply(self.init_weights) 40 | 41 | def _define_params(self): 42 | # define parameters in the model 43 | 44 | def forward(self, feed_dict): 45 | # generate prediction (ranking score according to tensors in feed_dict) 46 | item_id = feed_dict['item_id'] # [batch_size, -1] 47 | user_id = feed_dict['user_id'] # [batch_size] 48 | prediction = (...) 49 | out_dict = {'prediction': prediction.view(feed_dict['batch_size'], -1)} 50 | return out_dict 51 | 52 | class Dataset(GeneralModel.Dataset): 53 | # construct feed_dict for a single instance (called by __getitem__) 54 | # will be collated to a integrated feed dict for each batch 55 | def _get_feed_dict(self, index): 56 | feed_dict = super()._get_feed_dict(index) 57 | (...) 58 | return feed_dict 59 | ``` 60 | 61 | If the model definition is more complicated, you can inherit other functions in *BaseModel* (e.g. `loss`, `customize_parameters`) and *Dataset* (e.g. `_prepare`, `actions_before_epoch`), which needs deeper understandings about [BaseModel.py](https://github.com/THUwangcy/ReChorus/tree/master/src/models/BaseModel.py) and [BaseRunner.py](https://github.com/THUwangcy/ReChorus/tree/master/src/helpers/BaseRunner.py). You can also implement a new runner class to accommodate different experimental settings. 62 | -------------------------------------------------------------------------------- /src/exp.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import os 4 | import subprocess 5 | import pandas as pd 6 | import argparse 7 | import re 8 | import traceback 9 | import numpy as np 10 | from typing import List 11 | 12 | 13 | # Repeat experiments with different random seeds and save results to csv 14 | # Example: python exp.py --in_f run.sh --out_f exp.csv --n 5 15 | 16 | def parse_args(): 17 | parser = argparse.ArgumentParser(description="Run") 18 | parser.add_argument('--log_dir', nargs='?', default='../log/', 19 | help='Log save dir.') 20 | parser.add_argument('--cmd_dir', nargs='?', default='./', 21 | help='Command dir.') 22 | parser.add_argument('--in_f', nargs='?', default='run.sh', 23 | help='Input commands.') 24 | parser.add_argument('--out_f', nargs='?', default='exp.csv', 25 | help='Output csv.') 26 | parser.add_argument('--base_seed', type=int, default=0, 27 | help='Random seed at the beginning.') 28 | parser.add_argument('--n', type=int, default=5, 29 | help='Repeat times of each command.') 30 | parser.add_argument('--skip', type=int, default=0, 31 | help='skip number.') 32 | parser.add_argument('--gpu', type=str, default='0', 33 | help='Set CUDA_VISIBLE_DEVICES') 34 | return parser.parse_args() 35 | 36 | 37 | def find_info(result: List[str]) -> dict: 38 | info = dict() 39 | prefix = '' 40 | for line in result: 41 | if line.startswith(prefix + 'Best Iter(dev)'): 42 | line = line.replace(' ', '') 43 | p = re.compile('BestIter\(dev\)=(\d*)') 44 | info['Best Iter'] = p.search(line).group(1) 45 | p = re.compile('\[([\d\.]+)s\]') 46 | info['Time'] = p.search(line).group(1) 47 | elif line.startswith(prefix + 'Test After Training:'): 48 | p = re.compile('\(([\w@:\.\d,]+)\)') 49 | info['Test'] = p.search(line).group(1) 50 | return info 51 | 52 | 53 | def main(): 54 | args = parse_args() 55 | columns = ['Model', 'Test', 'Best Iter', 'Time', 'Seed', 'Run CMD'] 56 | skip = args.skip 57 | 58 | df = pd.DataFrame(columns=columns) 59 | # Read existing result record file 60 | if os.path.isfile(os.path.join(args.log_dir, args.out_f)): 61 | df = pd.read_csv(os.path.join(args.log_dir, args.out_f)) 62 | if (df.columns.to_numpy() != np.array(columns)).sum(): 63 | overwrite = input('Warning! There exists a file also called %s with different format, overwrite it? y/n (default n)'%(args.out_f)) 64 | if overwrite.lower().startswith('y'): 65 | print("Overwrite %s"%(args.out_f)) 66 | df = pd.DataFrame(columns=columns) 67 | else: 68 | exit(1) 69 | 70 | if not os.path.exists(args.log_dir): 71 | os.makedirs(args.log_dir) 72 | in_f = open(os.path.join(args.cmd_dir, args.in_f), 'r') 73 | lines = in_f.readlines() 74 | 75 | # Iterate commands 76 | for cmd in lines: 77 | cmd = cmd.strip() 78 | if cmd == '' or cmd.startswith('#') or cmd.startswith('export'): 79 | continue 80 | p = re.compile('--model_name (\w+)') 81 | model_name = p.search(cmd).group(1) 82 | 83 | # Repeat experiments 84 | for i in range(args.base_seed, args.base_seed + args.n): 85 | try: 86 | command = cmd 87 | if command.find(' --random_seed') == -1: 88 | command += ' --random_seed ' + str(i) 89 | if command.find(' --gpu ') == -1: 90 | command += ' --gpu ' + args.gpu 91 | if '${random_seed}' in command: 92 | command = command.replace('${random_seed}', str(i)) 93 | print(command) 94 | if skip > 0: 95 | skip -= 1 96 | continue 97 | result = subprocess.check_output(command, shell=True) 98 | result = result.decode('utf-8') 99 | result = [line.strip() for line in result.split(os.linesep)] 100 | info = find_info(result) 101 | info['Seed'] = str(i) 102 | info['Run CMD'] = command 103 | if args.n == 1: 104 | info['Model'] = model_name 105 | row = [info[c] if c in info else '' for c in columns] 106 | df.loc[len(df)] = row 107 | df.to_csv(os.path.join(args.log_dir, args.out_f), index=False) 108 | print(df[columns[:5]]) 109 | except Exception as e: 110 | traceback.print_exc() 111 | continue 112 | 113 | # Average results 114 | if args.n > 1: 115 | info = {'Model': model_name} 116 | tests = df['Test'].tolist()[-args.n:] 117 | tests_tuple =[[(m.split(':')[0],float(m.split(':')[1])) for m in t.split(',')] for t in tests] 118 | # avgs = ['{}:{:<.4f}'.format(np.average([t[mi] for t in tests])) for mi in range(len(tests[0]))] 119 | avgs = ['{}:{:<.4f}'.format(tests_tuple[0][mi][0],np.average([t[mi][1] for t in tests_tuple])) 120 | for mi in range(len(tests_tuple[0]))] 121 | info['Test'] = ','.join(avgs) 122 | epoch_avgs = np.mean([int(x) for x in df['Best Iter'].tolist()[-args.n:]]) 123 | info['Best Iter'] = "%.1f"%(epoch_avgs) 124 | row = [info[c] if c in info else '' for c in columns] 125 | df.loc[len(df)] = row 126 | print(df[columns[:5]]) 127 | for i in range(3): 128 | row = [''] * len(columns) 129 | df.loc[len(df)] = row 130 | 131 | # Save results 132 | df.to_csv(os.path.join(args.log_dir, args.out_f), index=False) 133 | 134 | 135 | if __name__ == '__main__': 136 | main() 137 | -------------------------------------------------------------------------------- /src/helpers/BUIRRunner.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import os 4 | import gc 5 | import torch 6 | import torch.nn as nn 7 | import logging 8 | import numpy as np 9 | from time import time 10 | from tqdm import tqdm 11 | from torch.utils.data import DataLoader 12 | 13 | from utils import utils 14 | from models.BaseModel import BaseModel 15 | from helpers.BaseRunner import BaseRunner 16 | 17 | 18 | class BUIRRunner(BaseRunner): 19 | def fit(self, dataset: BaseModel.Dataset, epoch=-1) -> float: 20 | model = dataset.model 21 | if model.optimizer is None: 22 | model.optimizer = self._build_optimizer(model) 23 | dataset.actions_before_epoch() # must sample before multi thread start 24 | 25 | model.train() 26 | loss_lst = list() 27 | dl = DataLoader(dataset, batch_size=self.batch_size, shuffle=True, num_workers=self.num_workers, 28 | collate_fn=dataset.collate_batch, pin_memory=self.pin_memory) 29 | for batch in tqdm(dl, leave=False, desc='Epoch {:<3}'.format(epoch), ncols=100, mininterval=1): 30 | batch = utils.batch_to_gpu(batch, model.device) 31 | model.optimizer.zero_grad() 32 | out_dict = model(batch) 33 | loss = model.loss(out_dict) 34 | loss.backward() 35 | model.optimizer.step() 36 | model._update_target() 37 | loss_lst.append(loss.detach().cpu().data.numpy()) 38 | return np.mean(loss_lst).item() 39 | -------------------------------------------------------------------------------- /src/helpers/BaseReader.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import os 4 | import pickle 5 | import argparse 6 | import logging 7 | import numpy as np 8 | import pandas as pd 9 | 10 | from utils import utils 11 | 12 | 13 | class BaseReader(object): 14 | @staticmethod 15 | def parse_data_args(parser): 16 | parser.add_argument('--path', type=str, default='data/', 17 | help='Input data dir.') 18 | parser.add_argument('--dataset', type=str, default='Grocery_and_Gourmet_Food', 19 | help='Choose a dataset.') 20 | parser.add_argument('--sep', type=str, default='\t', 21 | help='sep of csv file.') 22 | return parser 23 | 24 | def __init__(self, args): 25 | self.sep = args.sep 26 | self.prefix = args.path 27 | self.dataset = args.dataset 28 | self._read_data() 29 | 30 | self.train_clicked_set = dict() # store the clicked item set of each user in training set 31 | self.residual_clicked_set = dict() # store the residual clicked item set of each user 32 | for key in ['train', 'dev', 'test']: 33 | df = self.data_df[key] 34 | for uid, iid in zip(df['user_id'], df['item_id']): 35 | if uid not in self.train_clicked_set: 36 | self.train_clicked_set[uid] = set() 37 | self.residual_clicked_set[uid] = set() 38 | if key == 'train': 39 | self.train_clicked_set[uid].add(iid) 40 | else: 41 | self.residual_clicked_set[uid].add(iid) 42 | 43 | def _read_data(self): 44 | logging.info('Reading data from \"{}\", dataset = \"{}\" '.format(self.prefix, self.dataset)) 45 | self.data_df = dict() 46 | for key in ['train', 'dev', 'test']: 47 | self.data_df[key] = pd.read_csv(os.path.join(self.prefix, self.dataset, key + '.csv'), sep=self.sep).reset_index(drop=True).sort_values(by = ['user_id','time']) 48 | self.data_df[key] = utils.eval_list_columns(self.data_df[key]) 49 | 50 | logging.info('Counting dataset statistics...') 51 | key_columns = ['user_id','item_id','time'] 52 | if 'label' in self.data_df['train'].columns: # Add label for CTR prediction 53 | key_columns.append('label') 54 | self.all_df = pd.concat([self.data_df[key][key_columns] for key in ['train', 'dev', 'test']]) 55 | self.n_users, self.n_items = self.all_df['user_id'].max() + 1, self.all_df['item_id'].max() + 1 56 | for key in ['dev', 'test']: 57 | if 'neg_items' in self.data_df[key]: 58 | neg_items = np.array(self.data_df[key]['neg_items'].tolist()) 59 | assert (neg_items >= self.n_items).sum() == 0 # assert negative items don't include unseen ones 60 | logging.info('"# user": {}, "# item": {}, "# entry": {}'.format( 61 | self.n_users - 1, self.n_items - 1, len(self.all_df))) 62 | if 'label' in key_columns: 63 | positive_num = (self.all_df.label==1).sum() 64 | logging.info('"# positive interaction": {} ({:.1f}%)'.format( 65 | positive_num, positive_num/self.all_df.shape[0]*100)) 66 | 67 | -------------------------------------------------------------------------------- /src/helpers/CTRRunner.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import os 4 | import gc 5 | import torch 6 | import torch.nn as nn 7 | import logging 8 | import numpy as np 9 | from time import time 10 | from tqdm import tqdm 11 | from torch.utils.data import DataLoader 12 | from typing import Dict, List 13 | 14 | from utils import utils 15 | from models.BaseModel import BaseModel 16 | from helpers.BaseRunner import BaseRunner 17 | 18 | import sklearn.metrics as sk_metrics 19 | 20 | class CTRRunner(BaseRunner): 21 | 22 | @staticmethod 23 | def evaluate_method(predictions: np.ndarray,labels: np.ndarray, metrics: list) -> Dict[str, float]: 24 | """ 25 | :param predictions: An array of predictions for all samples 26 | :param labels: An array of labels for all samples (0 or 1) 27 | :param metrics: metric string list 28 | :return: a result dict, the keys are metrics 29 | """ 30 | evaluations = dict() 31 | for metric in metrics: 32 | if metric == 'ACC': 33 | evaluations[metric] = ((predictions>0.5).astype(int)==labels.astype(int)).mean() 34 | elif metric == 'AUC': 35 | evaluations[metric] = sk_metrics.roc_auc_score(labels,predictions) 36 | elif metric == 'F1_SCORE': 37 | evaluations[metric] = sk_metrics.f1_score(labels,(predictions>0.5).astype(int)) 38 | elif metric == 'LOG_LOSS': 39 | clip_predictions = np.clip(predictions, a_min=1e-7, a_max=1-1e-7) 40 | evaluations[metric] = -(np.log(clip_predictions)*labels+ np.log(1-clip_predictions)*(1-labels)).mean() 41 | else: 42 | raise ValueError('Undefined evaluation metric: {}.'.format(metric)) 43 | return evaluations 44 | 45 | def __init__(self, args): 46 | super().__init__(args) 47 | self.main_metric = self.metrics[0] if not len(args.main_metric) else self.main_metric 48 | 49 | def evaluate(self, dataset: BaseModel.Dataset, topks: list, metrics: list) -> Dict[str, float]: 50 | """ 51 | Evaluate the results for an input dataset. 52 | :return: result dict (key: metric) 53 | """ 54 | predictions, labels = self.predict(dataset) 55 | return self.evaluate_method(predictions, labels, metrics) 56 | 57 | def predict(self, dataset: BaseModel.Dataset, save_prediction: bool = False) -> np.ndarray: 58 | """ 59 | The returned prediction is a 1D-array corresponding to all samples, 60 | and ground truth labels are binary. 61 | """ 62 | dataset.model.eval() 63 | dataset.model.phase = 'eval' 64 | predictions, labels = list(), list() 65 | dl = DataLoader(dataset, batch_size=self.eval_batch_size, shuffle=False, num_workers=self.num_workers, 66 | collate_fn=dataset.collate_batch, pin_memory=self.pin_memory) 67 | for batch in tqdm(dl, leave=False, ncols=100, mininterval=1, desc='Predict'): 68 | if hasattr(dataset.model,'inference'): 69 | out_dict = dataset.model.inference(utils.batch_to_gpu(batch, dataset.model.device)) 70 | prediction, label = out_dict['prediction'], out_dict['label'] 71 | else: 72 | out_dict = dataset.model(utils.batch_to_gpu(batch, dataset.model.device)) 73 | prediction, label = out_dict['prediction'], out_dict['label'] 74 | predictions.extend(prediction.cpu().data.numpy()) 75 | labels.extend(label.cpu().data.numpy()) 76 | predictions = np.array(predictions) 77 | labels = np.array(labels) 78 | 79 | return predictions, labels 80 | -------------------------------------------------------------------------------- /src/helpers/ContextReader.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import logging 4 | import numpy as np 5 | import pandas as pd 6 | import os 7 | import sys 8 | 9 | from helpers.BaseReader import BaseReader 10 | 11 | ''' 12 | Reader for context information, including item, user, and situation context. 13 | ''' 14 | 15 | class ContextReader(BaseReader): 16 | @staticmethod 17 | def parse_data_args(parser): 18 | parser.add_argument('--include_item_features',type=int, default=0, 19 | help='Whether include item context features (0 or 1).') 20 | parser.add_argument('--include_user_features',type=int, default=0, 21 | help='Whether include user context features (0 or 1).') 22 | parser.add_argument('--include_situation_features',type=int, default=0, 23 | help='Whether include situation (i.e., dynamic context) features (0 or 1).') 24 | return BaseReader.parse_data_args(parser) 25 | 26 | def __init__(self, args): 27 | super().__init__(args) 28 | self.include_item_features = args.include_item_features 29 | self.include_user_features = args.include_user_features 30 | self.include_situation_features = args.include_situation_features 31 | self._load_ui_metadata() 32 | self._collect_context() 33 | 34 | def _load_ui_metadata(self): 35 | self.item_meta_df, self.user_meta_df = None, None 36 | item_meta_path = os.path.join(self.prefix, self.dataset, 'item_meta.csv') 37 | user_meta_path = os.path.join(self.prefix, self.dataset, 'user_meta.csv') 38 | if os.path.exists(item_meta_path) and self.include_item_features: 39 | self.item_meta_df = pd.read_csv(item_meta_path,sep=self.sep) 40 | self.item_feature_names = sorted([c for c in self.item_meta_df.columns if c[:2]=='i_']) 41 | else: 42 | self.item_feature_names = [] 43 | if os.path.exists(user_meta_path) and self.include_user_features: 44 | self.user_meta_df = pd.read_csv(user_meta_path,sep=self.sep) 45 | self.user_feature_names = sorted([c for c in self.user_meta_df.columns if c[:2]=='u_']) 46 | else: 47 | self.user_feature_names = [] 48 | if self.include_situation_features: 49 | self.situation_feature_names = sorted([c for c in self.data_df['train'].columns if c[:2]=='c_']) 50 | else: 51 | self.situation_feature_names = [] 52 | 53 | def _collect_context(self): 54 | logging.info('Collect context features...') 55 | id_columns = ['user_id','item_id'] 56 | self.item_features, self.user_features = None, None # dict 57 | self.feature_max = dict() 58 | for key in ['train', 'dev', 'test']: 59 | logging.info('Loading context for %s set...'%(key)) 60 | ids_df = self.data_df[key][id_columns] 61 | for f in id_columns: # get max value of each ID for embedding 62 | self.feature_max[f] = max(self.feature_max.get(f,0), int(ids_df[f].max())+1) 63 | # include situation features 64 | if self.include_situation_features and len(self.situation_feature_names): 65 | context_df = self.data_df[key][id_columns+['time']+self.situation_feature_names] 66 | for f in self.situation_feature_names: 67 | self.feature_max[f] = max(self.feature_max.get(f,0), int(context_df[f].max()) + 1 ) 68 | logging.info('#Situation Feautures: %d'%(context_df.shape[1]-3)) # except user id, item id, and user id 69 | del context_df 70 | # include item features 71 | if self.item_meta_df is not None and self.include_item_features: 72 | item_df = self.item_meta_df[['item_id']+self.item_feature_names] 73 | self.item_features = item_df.set_index('item_id').to_dict(orient='index') 74 | for f in self.item_feature_names: 75 | self.feature_max[f] = max( self.feature_max.get(f,0), int(item_df[f].max())+1 ) 76 | logging.info('# Item Features: %d'%(item_df.shape[1])) 77 | # include user features 78 | if self.user_meta_df is not None and self.include_user_features: 79 | user_df = self.user_meta_df[['user_id']+self.user_feature_names].set_index('user_id') 80 | self.user_features = user_df.to_dict(orient='index') 81 | for f in self.user_feature_names: 82 | self.feature_max[f] = max( self.feature_max.get(f,0), int(user_df[f].max())+1 ) 83 | logging.info('# User Features: %d'%(user_df.shape[1])) 84 | 85 | -------------------------------------------------------------------------------- /src/helpers/ContextSeqReader.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | ''' 3 | Jiayu Li 2023.05.20 4 | ''' 5 | 6 | import logging 7 | import pandas as pd 8 | import os 9 | import sys 10 | 11 | from helpers.ContextReader import ContextReader 12 | 13 | class ContextSeqReader(ContextReader): 14 | def __init__(self, args): 15 | super().__init__(args) 16 | self._append_his_info() 17 | 18 | def _append_his_info(self): 19 | """ 20 | Similar to SeqReader, but add situation context to each history interaction. 21 | self.user_his: store user history sequence [(i1,t1, {situation 1}), (i1,t2, {situation 2}), ...] 22 | """ 23 | logging.info('Appending history info with history context...') 24 | data_dfs = dict() 25 | for key in ['train','dev','test']: 26 | data_dfs[key] = self.data_df[key].copy() 27 | data_dfs[key]['phase'] = key 28 | sort_df = pd.concat([data_dfs[phase][['user_id','item_id','time','phase']+self.situation_feature_names] 29 | for phase in ['train','dev','test']]).sort_values(by=['time', 'user_id'], kind='mergesort') 30 | position = list() 31 | self.user_his = dict() # store the already seen sequence of each user 32 | situation_features = sort_df[self.situation_feature_names].to_numpy() 33 | for idx, (uid, iid, t) in enumerate(zip(sort_df['user_id'], sort_df['item_id'], sort_df['time'])): 34 | if uid not in self.user_his: 35 | self.user_his[uid] = list() 36 | position.append(len(self.user_his[uid])) 37 | self.user_his[uid].append((iid, t, situation_features[idx])) 38 | sort_df['position'] = position 39 | for key in ['train', 'dev', 'test']: 40 | self.data_df[key] = pd.merge( 41 | left=self.data_df[key], right=sort_df.drop(columns=['phase']+self.situation_feature_names), 42 | how='left', on=['user_id', 'item_id', 'time']) 43 | del sort_df -------------------------------------------------------------------------------- /src/helpers/ImpressionContextReader.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import logging 4 | import numpy as np 5 | import pandas as pd 6 | import os 7 | import sys 8 | 9 | from helpers.ImpressionReader import ImpressionReader 10 | from helpers.ContextReader import ContextReader 11 | from helpers.BaseReader import BaseReader 12 | from utils import utils 13 | 14 | class ImpressionContextReader(ImpressionReader, ContextReader): 15 | """ 16 | Impression-Context Reader reads impression data and add context information to it. 17 | """ 18 | 19 | @staticmethod 20 | def parse_data_args(parser): 21 | parser.add_argument('--include_item_features',type=int, default=0, 22 | help='Whether include item context features.') 23 | parser.add_argument('--include_user_features',type=int, default=0, 24 | help='Whether include user context features.') 25 | parser.add_argument('--include_context_features',type=int, default=0, 26 | help='Whether include dynamic context features.') 27 | return BaseReader.parse_data_args(parser) 28 | 29 | def __init__(self, args): 30 | self.sep = args.sep 31 | self.prefix = args.path 32 | self.dataset = args.dataset 33 | self._read_data() 34 | 35 | self.train_clicked_set = dict() # store the clicked item set of each user in training set 36 | self.residual_clicked_set = dict() # store the residual clicked item set of each user 37 | for key in ['train', 'dev', 'test']: 38 | df = self.data_df[key] 39 | for uid, iid in zip(df['user_id'], df['item_id']): 40 | if uid not in self.train_clicked_set: 41 | self.train_clicked_set[uid] = set() 42 | self.residual_clicked_set[uid] = set() 43 | if key == 'train': 44 | self.train_clicked_set[uid].add(iid) 45 | else: 46 | self.residual_clicked_set[uid].add(iid) 47 | self.include_item_features = args.include_item_features 48 | self.include_user_features = args.include_user_features 49 | self.include_context_features = args.include_context_features 50 | self._load_ui_metadata() 51 | self._collect_context() 52 | self._append_impression_info() 53 | -------------------------------------------------------------------------------- /src/helpers/ImpressionReader.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import logging 4 | import numpy as np 5 | import pandas as pd 6 | import os 7 | import sys 8 | 9 | from helpers.BaseReader import BaseReader 10 | from utils import utils 11 | 12 | class ImpressionReader(BaseReader): 13 | """ 14 | Impression Reader reads impression data. In each impression there are pre-defined unfixed number of positive items and negative items 15 | """ 16 | @staticmethod 17 | def parse_data_args(parser): 18 | parser.add_argument('--impression_idkey', type=str, default='time', 19 | help='The key for impression identification, [time, impression_id]') 20 | return BaseReader.parse_data_args(parser) 21 | 22 | def __init__(self, args): 23 | self.impression_idkey = args.impression_idkey 24 | super().__init__(args) 25 | self._append_impression_info() 26 | 27 | def _read_data(self): 28 | logging.info('Reading data from \"{}\", dataset = \"{}\" '.format(self.prefix, self.dataset)) 29 | self.data_df = dict() 30 | for key in ['train', 'dev', 'test']: 31 | self.data_df[key] = pd.read_csv(os.path.join(self.prefix, self.dataset, key + '.csv'), sep=self.sep).reset_index(drop=True).sort_values(by = ['user_id',self.impression_idkey]) 32 | self.data_df[key] = utils.eval_list_columns(self.data_df[key]) 33 | logging.info('Counting dataset statistics...') 34 | if self.impression_idkey == 'time': 35 | key_columns = ['user_id', 'item_id', 'time'] 36 | else: 37 | key_columns = ['user_id', 'item_id', 'time', self.impression_idkey] 38 | if 'label' in self.data_df['train'].columns: 39 | key_columns.append('label') 40 | else: 41 | raise KeyError('Impression data must have binary labels') 42 | self.all_df = pd.concat([self.data_df[key][key_columns] for key in ['train', 'dev', 'test']]) 43 | self.n_users, self.n_items = self.all_df['user_id'].max() + 1, self.all_df['item_id'].max() + 1 44 | # In impression data, negative item lists can have unseen items (i.e., items without click) 45 | logging.info('Update impression data -- "# user": {}, "# item": {}, "# entry": {}'.format( 46 | self.n_users - 1, self.n_items - 1, len(self.all_df))) 47 | if 'label' in key_columns: 48 | positive_num = (self.all_df.label==1).sum() 49 | logging.info('"# positive interaction": {} ({:.1f}%)'.format( 50 | positive_num, positive_num/self.all_df.shape[0]*100)) 51 | 52 | def _append_impression_info(self): # -> NoReturn: 53 | """ 54 | Merge all positive items of a request based on the timestamp/impression_idkey, and get column 'pos_items' for self.data_df 55 | Add impression info to data_df: neg_num, pos_num 56 | """ 57 | logging.info('Merging positive items by timestamp/impression_idkey...') 58 | # train,val,test 59 | mask = {'train':[],'dev':[],'test':[]} 60 | for key in self.data_df.keys(): 61 | df=self.data_df[key].copy() 62 | df['last_user_id'] = df['user_id'].shift(1) 63 | df['last_'+self.impression_idkey] = df[self.impression_idkey].shift(1) 64 | 65 | positive_items, negative_items = [], [] 66 | current_pos, current_neg = set(), set() 67 | for uid, last_uid, ipid, last_ipid, iid, label in \ 68 | df[['user_id','last_user_id',self.impression_idkey,'last_'+self.impression_idkey,'item_id','label']].to_numpy(): 69 | if uid == last_uid and ipid == last_ipid: 70 | positive_items.append([]) 71 | negative_items.append([]) 72 | mask[key].append(0) 73 | else: 74 | if len(current_pos): 75 | positive_items.append(list(current_pos)) 76 | negative_items.append(list(current_neg)) 77 | mask[key].append(1) 78 | else: 79 | if len(current_neg):#impression with only neg items are dropped 80 | positive_items.append([]) 81 | negative_items.append([]) 82 | mask[key].append(0) 83 | current_pos, current_neg = set(), set() 84 | if label: 85 | current_pos = current_pos.union(set([iid])) 86 | else: 87 | current_neg = current_neg.union(set([iid])) 88 | # last session 89 | if len(current_pos): 90 | positive_items.append(list(current_pos)) 91 | negative_items.append(list(current_neg)) 92 | mask[key].append(1) 93 | else: 94 | if len(current_neg):#impression with only neg items are dropped 95 | positive_items.append([]) 96 | negative_items.append([]) 97 | mask[key].append(0) 98 | self.data_df[key]['pos_items'] = positive_items 99 | self.data_df[key]['neg_items'] = negative_items 100 | self.data_df[key]=self.data_df[key][np.array(mask[key])==1] 101 | 102 | logging.info('Appending neg_num & pos_num...') 103 | 104 | neg_num_sum, pos_num_sum = 0,0 105 | for key in ['train', 'dev', 'test']: 106 | df = self.data_df[key] 107 | neg_num = list() 108 | pos_num = list() 109 | for neg_items in df['neg_items']: 110 | if 0 in neg_items: 111 | neg_num.append(neg_items.index(0)) 112 | else: 113 | neg_num.append(len(neg_items)) 114 | self.data_df[key]['neg_num']=neg_num 115 | for pos_items in df['pos_items']: 116 | if 0 in pos_items: 117 | pos_num.append(pos_items.index(0)) 118 | else: 119 | pos_num.append(len(pos_items)) 120 | self.data_df[key]['pos_num']=pos_num 121 | self.data_df[key] = self.data_df[key].loc[self.data_df[key].neg_num>0].reset_index(drop=True) # Retain sessions with negative data only 122 | neg_num_sum += sum(neg_num) 123 | pos_num_sum += sum(pos_num) 124 | neg_num_avg = neg_num_sum / sum([self.data_df[key].shape[0] for key in self.data_df]) 125 | pos_num_avg = pos_num_sum / sum([self.data_df[key].shape[0] for key in self.data_df]) 126 | 127 | logging.info('train, dev, test request num: '+str(len(self.data_df['train']))+' '+str(len(self.data_df['dev']))+' '+str(len(self.data_df['test']))) 128 | logging.info("Average positive items / impression = %.3f, negative items / impression = %.3f"%( 129 | pos_num_avg,neg_num_avg)) -------------------------------------------------------------------------------- /src/helpers/ImpressionSeqReader.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import logging 4 | import numpy as np 5 | import pandas as pd 6 | import os 7 | import sys 8 | 9 | from helpers.ImpressionReader import ImpressionReader 10 | from utils import utils 11 | 12 | class ImpressionSeqReader(ImpressionReader): 13 | 14 | def __init__(self, args): 15 | super().__init__(args) 16 | self._append_his_info() 17 | 18 | def _append_his_info(self): 19 | """ 20 | self.user_his: store both positive and negative user history sequences 21 | pos: [(i1,t1), (i1,t2), ...]; 22 | neg: [(in1,tn1), (in2,tn2),...] 23 | add the 'position' of each interaction in user_his to data_df 24 | """ 25 | logging.info('Appending history info with corresponding impressions...') 26 | data_dfs = dict() 27 | for key in ['train','dev','test']: 28 | data_dfs[key] = self.data_df[key].copy() 29 | data_dfs[key]['phase'] = key 30 | if self.impression_idkey == 'time': 31 | key_columns = ['user_id', 'pos_items', 'neg_items', 'time', 'phase'] 32 | sort_columns = ['user_id', 'time'] 33 | else: 34 | key_columns = ['user_id', 'pos_items', 'neg_items', 'time', 'phase', self.impression_idkey] 35 | sort_columns = ['user_id', self.impression_idkey, 'time'] 36 | sort_df = pd.concat([data_dfs[phase][key_columns] 37 | for phase in ['train','dev','test']]).sort_values(by=sort_columns, kind='mergesort') 38 | position = list() 39 | neg_position = list() 40 | self.user_his = dict() # store the already seen sequence of each user 41 | for idx, (uid, pids, nids, t) in enumerate(zip(sort_df['user_id'], sort_df['pos_items'], sort_df['neg_items'], sort_df['time'])): 42 | if uid not in self.user_his: 43 | self.user_his[uid] = {'pos':list(),'neg':list()} 44 | position.append(len(self.user_his[uid]['pos'])) 45 | neg_position.append(len(self.user_his[uid]['neg'])) 46 | for pid in pids: 47 | self.user_his[uid]['pos'].append((pid, t)) 48 | for nid in nids: 49 | self.user_his[uid]['neg'].append((nid, t)) 50 | sort_df['position'] = position 51 | sort_df['neg_position'] = neg_position 52 | for key in ['train', 'dev', 'test']: 53 | self.data_df[key] = pd.merge( 54 | left=self.data_df[key], right=sort_df.drop(columns=['phase','pos_items','neg_items']), 55 | how='left', on=['user_id', self.impression_idkey]) 56 | del sort_df -------------------------------------------------------------------------------- /src/helpers/KDAReader.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import os 4 | import pickle 5 | import argparse 6 | import logging 7 | import numpy as np 8 | import pandas as pd 9 | from tqdm import tqdm 10 | 11 | from helpers.KGReader import KGReader 12 | 13 | 14 | """ Data Reading for KDA """ 15 | class KDAReader(KGReader): 16 | @staticmethod 17 | def parse_data_args(parser): 18 | parser.add_argument('--t_scalar', type=int, default=60, 19 | help='Time interval scalar.') 20 | parser.add_argument('--n_dft', type=int, default=64, 21 | help='The point of DFT.') 22 | parser.add_argument('--freq_rand', type=int, default=0, 23 | help='Whether randomly initialize parameters in frequency domain.') 24 | return KGReader.parse_data_args(parser) 25 | 26 | @staticmethod 27 | def dft(x: list, n_dft=-1) -> np.ndarray: 28 | if n_dft <= 0: 29 | n_dft = 2 ** (int(np.log2(len(x))) + 1) 30 | freq_x = np.fft.fft(x, n_dft) 31 | return 2 * freq_x[: n_dft // 2 + 1] # fold negative frequencies 32 | 33 | @staticmethod 34 | def norm_time(a: list, t_scalar: int) -> np.ndarray: 35 | norm_t = np.log2(np.array(a) / t_scalar + 1e-6) 36 | norm_t = np.maximum(norm_t, 0) 37 | return norm_t 38 | 39 | def __init__(self, args): 40 | super().__init__(args) 41 | self.t_scalar = args.t_scalar 42 | self.n_dft = args.n_dft 43 | self.freq_rand = args.freq_rand 44 | self.regenerate = args.regenerate 45 | self.interval_file = os.path.join(self.prefix, self.dataset, 'interval.pkl') 46 | 47 | self.freq_x = np.empty((self.n_relations, self.n_dft // 2 + 1), dtype=complex) 48 | if not self.freq_rand: 49 | self._time_interval_cnt() # ! May need a lot of time 50 | self._cal_freq_x() 51 | 52 | # Calculate time intervals of relational neighbors for each relation type (include a virtual relation) 53 | def _time_interval_cnt(self): 54 | if os.path.exists(self.interval_file) and not self.regenerate: 55 | self.interval_dict = pickle.load(open(self.interval_file, 'rb')) 56 | return 57 | 58 | self.interval_dict = {'virtual': list()} 59 | for relation_type in self.relations: 60 | self.interval_dict[relation_type] = list() 61 | 62 | merge_df = pd.merge(self.all_df, self.item_meta_df, how='left', on='item_id') 63 | gb_user = merge_df.groupby('user_id') 64 | for user, user_df in tqdm(gb_user, leave=False, ncols=100, mininterval=1, desc='Count Intervals'): 65 | # Virtual item-item relation 66 | times, iids = user_df['time'].values, user_df['item_id'].values 67 | delta_t = [t for t in (times[1:] - times[:-1]) if t > 0] 68 | self.interval_dict['virtual'].extend(delta_t) 69 | # Attribute based relations 70 | for attr in self.attr_relations: 71 | for val, df in user_df.groupby(attr): 72 | delta_t = [t for t in (df['time'].values[1:] - df['time'].values[:-1]) if t > 0] 73 | self.interval_dict[attr].extend(delta_t) 74 | # Natural item relations 75 | for r_idx, relation in enumerate(self.item_relations): 76 | for target_idx in range(1, len(iids))[::-1]: # traverse tail item back-to-front in user history 77 | target_i, target_t = iids[target_idx], times[target_idx] 78 | for source_idx in range(target_idx)[::-1]: # look forward from the tail item 79 | source_i, source_t = iids[source_idx], times[source_idx] 80 | delta_t = target_t - source_t 81 | if delta_t > 0 and (source_i, r_idx + 1, target_i) in self.triplet_set: 82 | self.interval_dict[relation].append(delta_t) 83 | break 84 | 85 | pickle.dump(self.interval_dict, open(self.interval_file, 'wb')) 86 | 87 | # Apply DFT on time interval distributions to get initial frequency representations 88 | def _cal_freq_x(self): 89 | distributions = list() 90 | # normalized time interval distributions 91 | for col in ['virtual'] + self.relations: 92 | intervals = self.norm_time(self.interval_dict[col], self.t_scalar) 93 | bin_num = int(max(intervals)) + 1 94 | ns = np.zeros(bin_num) 95 | for inter in intervals: 96 | ns[int(inter)] += 1 97 | distributions.append(ns / max(ns)) 98 | min_dft = 2 ** (int(np.log2(bin_num) + 1)) 99 | if self.n_dft < min_dft: 100 | self.n_dft = min_dft 101 | # DFT 102 | for i, dist in enumerate(distributions): 103 | dft_res = self.dft(dist, self.n_dft) 104 | self.freq_x[i] = dft_res 105 | 106 | del self.interval_dict 107 | -------------------------------------------------------------------------------- /src/helpers/KGReader.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import os 4 | import pickle 5 | import argparse 6 | import logging 7 | import numpy as np 8 | import pandas as pd 9 | 10 | 11 | from helpers.SeqReader import SeqReader 12 | from utils import utils 13 | 14 | 15 | class KGReader(SeqReader): 16 | @staticmethod 17 | def parse_data_args(parser): 18 | parser.add_argument('--include_attr', type=int, default=0, 19 | help='Whether include attribute-based relations.') 20 | return SeqReader.parse_data_args(parser) 21 | 22 | def __init__(self, args): 23 | super().__init__(args) 24 | self.include_attr = args.include_attr 25 | item_meta_path = os.path.join(self.prefix, self.dataset, 'item_meta.csv') 26 | self.item_meta_df = pd.read_csv(item_meta_path, sep=self.sep) 27 | self.item_meta_df = utils.eval_list_columns(self.item_meta_df) 28 | 29 | self._construct_kg() 30 | 31 | def _construct_kg(self): 32 | logging.info('Constructing relation triplets...') 33 | 34 | self.triplet_set = set() 35 | heads, relations, tails = [], [], [] 36 | 37 | self.item_relations = [r for r in self.item_meta_df.columns if r.startswith('r_')] 38 | for idx in range(len(self.item_meta_df)): 39 | head_item = self.item_meta_df['item_id'].values[idx] 40 | for r_idx, r in enumerate(self.item_relations): 41 | for tail_item in self.item_meta_df[r].values[idx]: 42 | heads.append(head_item) 43 | tails.append(tail_item) 44 | relations.append(r_idx + 1) # idx 0 is reserved to be a virtual relation between items 45 | self.triplet_set.add((head_item, r_idx + 1, tail_item)) 46 | logging.info('Item-item relations:' + str(self.item_relations)) 47 | 48 | self.attr_relations = list() 49 | if self.include_attr: 50 | self.attr_relations = [r for r in self.item_meta_df.columns if r.startswith('i_')] 51 | self.attr_max, self.share_attr_dict = list(), dict() 52 | for r_idx, attr in enumerate(self.attr_relations): 53 | base = self.n_items + np.sum(self.attr_max) # base index of attribute entities 54 | relation_idx = len(self.item_relations) + r_idx + 1 # index of the relation type 55 | for item, val in zip(self.item_meta_df['item_id'], self.item_meta_df[attr]): 56 | if val != 0: # the attribute is not NaN 57 | heads.append(item) 58 | tails.append(int(val + base)) 59 | relations.append(relation_idx) 60 | self.triplet_set.add((item, relation_idx, int(val + base))) 61 | for val, val_df in self.item_meta_df.groupby(attr): 62 | self.share_attr_dict[int(val + base)] = val_df['item_id'].tolist() 63 | self.attr_max.append(self.item_meta_df[attr].max() + 1) 64 | logging.info('Attribute-based relations:' + str(self.attr_relations)) 65 | 66 | self.relations = self.item_relations + self.attr_relations 67 | self.relation_df = pd.DataFrame() 68 | self.relation_df['head'] = heads 69 | self.relation_df['relation'] = relations 70 | self.relation_df['tail'] = tails 71 | self.n_relations = len(self.relations) + 1 72 | self.n_entities = pd.concat((self.relation_df['head'], self.relation_df['tail'])).max() + 1 73 | logging.info('"# relation": {}, "# triplet": {}'.format(self.n_relations, len(self.relation_df))) 74 | 75 | 76 | if __name__ == '__main__': 77 | logging.basicConfig(level=logging.INFO) 78 | parser = argparse.ArgumentParser() 79 | parser = KGReader.parse_data_args(parser) 80 | args, extras = parser.parse_known_args() 81 | 82 | args.path = '../../data/' 83 | corpus = KGReader(args) 84 | 85 | corpus_path = os.path.join(args.path, args.dataset, 'KGReader.pkl') 86 | logging.info('Save corpus to {}'.format(corpus_path)) 87 | pickle.dump(corpus, open(corpus_path, 'wb')) 88 | -------------------------------------------------------------------------------- /src/helpers/SeqReader.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import logging 4 | import pandas as pd 5 | 6 | from helpers.BaseReader import BaseReader 7 | 8 | 9 | class SeqReader(BaseReader): 10 | def __init__(self, args): 11 | super().__init__(args) 12 | self._append_his_info() 13 | 14 | def _append_his_info(self): 15 | """ 16 | self.user_his: store user history sequence [(i1,t1), (i1,t2), ...] 17 | add the 'position' of each interaction in user_his to data_df 18 | """ 19 | logging.info('Appending history info...') 20 | sort_df = self.all_df.sort_values(by=['time', 'user_id'], kind='mergesort') 21 | position = list() 22 | self.user_his = dict() # store the already seen sequence of each user 23 | for uid, iid, t in zip(sort_df['user_id'], sort_df['item_id'], sort_df['time']): 24 | if uid not in self.user_his: 25 | self.user_his[uid] = list() 26 | position.append(len(self.user_his[uid])) 27 | self.user_his[uid].append((iid, t)) 28 | sort_df['position'] = position 29 | for key in ['train', 'dev', 'test']: 30 | self.data_df[key] = pd.merge( 31 | left=self.data_df[key], right=sort_df, how='left', 32 | on=['user_id', 'item_id', 'time']) 33 | del sort_df -------------------------------------------------------------------------------- /src/helpers/__init__.py: -------------------------------------------------------------------------------- 1 | from os.path import dirname, basename, isfile, join 2 | import glob 3 | 4 | modules = glob.glob(join(dirname(__file__), "*.py")) 5 | __all__ = [ 6 | basename(f)[:-3] for f in modules if isfile(f) and not f.endswith('__init__.py') 7 | ] 8 | -------------------------------------------------------------------------------- /src/models/BaseContextModel.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import torch 4 | import logging 5 | import numpy as np 6 | from tqdm import tqdm 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | from torch.nn.utils.rnn import pad_sequence 10 | from typing import List 11 | 12 | from utils import utils 13 | from models.BaseModel import * 14 | 15 | def get_context_feature(feed_dict, index, corpus, data): 16 | """ 17 | Get context features for the feed_dict, including user, item, and situation context 18 | """ 19 | for c in corpus.user_feature_names: 20 | feed_dict[c] = corpus.user_features[feed_dict['user_id']][c] 21 | for c in corpus.situation_feature_names: 22 | feed_dict[c] = data[c][index] 23 | for c in corpus.item_feature_names: 24 | if type(feed_dict['item_id']) in [int, np.int32, np.int64]: # for a single item 25 | feed_dict[c] = corpus.item_features[feed_dict['item_id']][c] 26 | else: # for item list 27 | feed_dict[c] = np.array([corpus.item_features[iid][c] for iid in feed_dict['item_id']]) 28 | return feed_dict 29 | 30 | class ContextModel(GeneralModel): 31 | # context model for top-k recommendation tasks 32 | reader = 'ContextReader' 33 | 34 | @staticmethod 35 | def parse_model_args(parser): 36 | parser.add_argument('--loss_n',type=str,default='BPR', 37 | help='Type of loss functions.') 38 | return GeneralModel.parse_model_args(parser) 39 | 40 | def __init__(self, args, corpus): 41 | super().__init__(args, corpus) 42 | self.loss_n = args.loss_n 43 | self.context_features = corpus.user_feature_names + corpus.item_feature_names + corpus.situation_feature_names\ 44 | + ['user_id','item_id'] 45 | self.feature_max = corpus.feature_max 46 | 47 | def loss(self, out_dict: dict): 48 | """ 49 | utilize BPR loss (same as general models) or BCE loss (same as CTR tasks) 50 | """ 51 | if self.loss_n == 'BPR': 52 | loss = super().loss(out_dict) 53 | elif self.loss_n == 'BCE': 54 | predictions = out_dict['prediction'].sigmoid() 55 | pos_pred, neg_pred = predictions[:, 0], predictions[:, 1:] 56 | loss = - (pos_pred.log() + (1-neg_pred).log().sum(dim=1)).mean() 57 | else: 58 | raise ValueError('Undefined loss function: {}'.format(self.loss_n)) 59 | if torch.isnan(loss) or torch.isinf(loss): 60 | print('Error!') 61 | return loss 62 | 63 | class Dataset(GeneralModel.Dataset): 64 | def __init__(self, model, corpus, phase): 65 | super().__init__(model, corpus, phase) 66 | 67 | def _get_feed_dict(self, index): 68 | feed_dict = super()._get_feed_dict(index) 69 | feed_dict = get_context_feature(feed_dict, index, self.corpus, self.data) 70 | return feed_dict 71 | 72 | 73 | class ContextCTRModel(CTRModel): 74 | # context model for CTR prediction tasks 75 | reader = 'ContextReader' 76 | 77 | def __init__(self, args, corpus): 78 | super().__init__(args, corpus) 79 | self.context_features = corpus.user_feature_names + corpus.item_feature_names + corpus.situation_feature_names\ 80 | + ['user_id','item_id'] 81 | self.feature_max = corpus.feature_max 82 | 83 | class Dataset(CTRModel.Dataset): 84 | def _get_feed_dict(self, index): 85 | feed_dict = super()._get_feed_dict(index) 86 | feed_dict = get_context_feature(feed_dict, index, self.corpus, self.data) 87 | return feed_dict 88 | 89 | class ContextSeqModel(ContextModel): 90 | reader='ContextSeqReader' 91 | 92 | @staticmethod 93 | def parse_model_args(parser): 94 | parser.add_argument('--history_max', type=int, default=20, 95 | help='Maximum length of history.') 96 | parser.add_argument('--add_historical_situations',type=int,default=0, 97 | help='Whether to add historical situation context as sequence.') 98 | return ContextModel.parse_model_args(parser) 99 | 100 | def __init__(self, args, corpus): 101 | super().__init__(args, corpus) 102 | self.history_max = args.history_max 103 | self.add_historical_situations = args.add_historical_situations 104 | 105 | class Dataset(SequentialModel.Dataset): 106 | def __init__(self, model, corpus, phase): 107 | super().__init__(model, corpus, phase) 108 | 109 | def _get_feed_dict(self, index): 110 | # get item features, user features, and context features separately 111 | feed_dict = super()._get_feed_dict(index) 112 | pos = self.data['position'][index] 113 | user_seq = self.corpus.user_his[feed_dict['user_id']][:pos] 114 | if self.model.history_max > 0: 115 | user_seq = user_seq[-self.model.history_max:] 116 | feed_dict = get_context_feature(feed_dict, index, self.corpus, self.data) 117 | for c in self.corpus.item_feature_names: # get historical item context features 118 | feed_dict['history_'+c] = np.array([self.corpus.item_features[iid][c] for iid in feed_dict['history_items']]) 119 | if self.model.add_historical_situations: # get historical situation context features 120 | for idx,c in enumerate(self.corpus.situation_feature_names): 121 | feed_dict['history_'+c] = np.array([inter[-1][idx] for inter in user_seq]) 122 | feed_dict['history_item_id'] = feed_dict['history_items'] 123 | feed_dict.pop('history_items') 124 | return feed_dict 125 | 126 | class ContextSeqCTRModel(ContextCTRModel): 127 | reader = 'ContextSeqReader' 128 | 129 | @staticmethod 130 | def parse_model_args(parser): 131 | parser.add_argument('--history_max', type=int, default=20, 132 | help='Maximum length of history.') 133 | parser.add_argument('--add_historical_situations',type=int,default=0, 134 | help='Whether to add historical situation context as sequence.') 135 | return ContextCTRModel.parse_model_args(parser) 136 | 137 | def __init__(self, args, corpus): 138 | super().__init__(args, corpus) 139 | self.history_max = args.history_max 140 | self.add_historical_situations = args.add_historical_situations 141 | 142 | class Dataset(ContextCTRModel.Dataset): 143 | def __init__(self, model, corpus, phase): 144 | super().__init__(model, corpus, phase) 145 | idx_select = np.array(self.data['position']) > 0 # history length must be non-zero 146 | for key in self.data: 147 | self.data[key] = np.array(self.data[key])[idx_select] 148 | 149 | def _get_feed_dict(self, index): 150 | feed_dict = super()._get_feed_dict(index) 151 | pos = self.data['position'][index] 152 | user_seq = self.corpus.user_his[feed_dict['user_id']][:pos] 153 | if self.model.history_max > 0: 154 | user_seq = user_seq[-self.model.history_max:] 155 | # feed_dict = get_context_feature(feed_dict, index, self.corpus, self.data) 156 | feed_dict['history_items'] = np.array([x[0] for x in user_seq]) 157 | feed_dict['history_times'] = np.array([x[1] for x in user_seq]) 158 | feed_dict['lengths'] = len(feed_dict['history_items']) 159 | for c in self.corpus.item_feature_names: # get historical item context features 160 | feed_dict['history_'+c] = np.array([self.corpus.item_features[iid][c] for iid in feed_dict['history_items']]) 161 | if self.model.add_historical_situations: # get historical situation context features 162 | for idx,c in enumerate(self.corpus.situation_feature_names): 163 | feed_dict['history_'+c] = np.array([inter[-1][idx] for inter in user_seq]) 164 | feed_dict['history_item_id'] = feed_dict['history_items'] 165 | feed_dict.pop('history_items') 166 | return feed_dict 167 | -------------------------------------------------------------------------------- /src/models/BaseRerankerModel.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import torch.nn.functional as F 4 | from typing import List 5 | import yaml 6 | import copy 7 | 8 | from models.BaseModel import * 9 | from models.BaseImpressionModel import * 10 | from models.general import * 11 | from models.sequential import * 12 | from models.developing import * 13 | 14 | 15 | class RerankModel(ImpressionModel): 16 | reader='ImpressionReader' 17 | runner='ImpressionRunner' 18 | extra_log_args = ['tuneranker'] 19 | 20 | @staticmethod 21 | def parse_model_args(parser): 22 | parser.add_argument('--ranker_name', type=str, default='BPRMF_test', 23 | help='Base ranker') 24 | parser.add_argument('--ranker_config_file', type=str, default='BPRMF_test.yaml', 25 | help='Base ranker config file') 26 | parser.add_argument('--ranker_model_file', type=str, default='BPRMF_test__MINDCTR__2__lr=0.001__l2=0.0__emb_size=64__batch_size=256.pt', 27 | help='Base ranker model file') 28 | parser.add_argument('--tuneranker', type=int, default=0, 29 | help='if 1, continue to train ranker') 30 | return ImpressionModel.parse_model_args(parser) 31 | 32 | def __init__(self, args, corpus): 33 | super().__init__(args, corpus) 34 | self.ranker_name = args.ranker_name 35 | self.ranker_config = args.ranker_config_file 36 | self.ranker_model = args.ranker_model_file 37 | self.tuneranker = args.tuneranker 38 | self.load_ranker(args, corpus) 39 | 40 | def load_ranker(self, args, corpus): 41 | corpus = corpus 42 | # config_path = './{}'.format(self.ranker_config) 43 | # model_path = './{}'.format(self.ranker_model) 44 | config_path = './model/{}Impression/{}'.format(self.ranker_name,self.ranker_config) 45 | model_path = './model/{}Impression/{}'.format(self.ranker_name,self.ranker_model) 46 | #read ranker config 47 | ranker_config_dict = dict() 48 | with open(config_path, "r", encoding="utf-8") as f: 49 | ranker_config_dict.update( 50 | yaml.load(f.read(), Loader=yaml.FullLoader) 51 | ) 52 | #load ranker from model and get self.ranker_emb_size 53 | ranker_args = copy.deepcopy(args) 54 | for k, v in ranker_config_dict.items(): 55 | if k != 'history_max': 56 | setattr(ranker_args, k, v) 57 | model_name = eval('{0}.{0}Impression'.format(self.ranker_name)) 58 | self.ranker = model_name(ranker_args, corpus) 59 | self.ranker.device = ranker_args.device 60 | self.ranker.apply(self.ranker.init_weights) 61 | self.ranker.to(self.device) 62 | self.ranker_emb_size = ranker_args.emb_size 63 | self.ranker.load_model(model_path) 64 | if not self.tuneranker: 65 | for param in self.ranker.parameters(): 66 | param.requires_grad = False 67 | 68 | class Dataset(ImpressionModel.Dataset): 69 | # Collate a batch according to the list of feed dicts 70 | def collate_batch(self, feed_dicts: List[dict]) -> dict: # feed_dicts are a batch of dicts 71 | feed_dict = super().collate_batch(feed_dicts) 72 | feed_dict['batch_size'] = len(feed_dicts) 73 | predict_dict = self.model.ranker(utils.batch_to_gpu(feed_dict, self.model.device)) # pos+pad+neg+pad 74 | feed_dict['scores'] = predict_dict['prediction'] # [batch(or num_sequence),n_candidate] 75 | pos_mask = torch.arange(0, self.model.train_max_pos_item, device = self.model.device).type_as(feed_dict['pos_num']).unsqueeze(0).expand(feed_dict['batch_size'], self.model.train_max_pos_item).lt(feed_dict['pos_num'].unsqueeze(1)) 76 | neg_mask = torch.arange(0, self.model.train_max_neg_item, device = self.model.device).type_as(feed_dict['neg_num']).unsqueeze(0).expand(feed_dict['batch_size'], self.model.train_max_neg_item).lt(feed_dict['neg_num'].unsqueeze(1)) 77 | all_mask = torch.cat([pos_mask, neg_mask],dim = 1) 78 | feed_dict['padding_mask'] = ~all_mask 79 | feed_dict['scores'] = torch.where(all_mask == 1, feed_dict['scores'],-np.inf * torch.ones_like(feed_dict['scores'])) 80 | _,temp = feed_dict['scores'].sort(dim = 1, descending = True) 81 | _,feed_dict['position'] = temp.sort(dim = 1) 82 | feed_dict['u_v'] = predict_dict['u_v'] # [batch(or num_sequence),ranker_embedding_len] 83 | feed_dict['i_v'] = predict_dict['i_v'] # [batch(or num_sequence),ranker_embedding_len] 84 | return feed_dict 85 | 86 | class RerankSeqModel(RerankModel): 87 | reader='ImpressionSeqReader' 88 | runner='ImpressionRunner' 89 | extra_log_args = ['tuneranker'] 90 | 91 | @staticmethod 92 | def parse_model_args(parser): 93 | parser.add_argument('--history_max', type=int, default=20, 94 | help='Maximum length of history.') 95 | parser.add_argument('--ranker_name', type=str, default='SASRec_test', 96 | help='Base ranker') 97 | parser.add_argument('--ranker_config_file', type=str, default='SASRec_test.yaml', 98 | help='Base ranker config file') 99 | parser.add_argument('--ranker_model_file', type=str, default='SASRec_test__MINDCTR__1__lr=0.0005__l2=0.0__emb_size=64__num_layers=3__num_heads=1.pt', 100 | help='Base ranker model file') 101 | parser.add_argument('--tuneranker', type=int, default=0, 102 | help='if 1, continue to train ranker') 103 | return ImpressionModel.parse_model_args(parser) 104 | 105 | def __init__(self, args, corpus): 106 | super().__init__(args, corpus) 107 | self.history_max = args.history_max 108 | 109 | class Dataset(ImpressionSeqModel.Dataset): 110 | # Collate a batch according to the list of feed dicts 111 | def collate_batch(self, feed_dicts: List[dict]) -> dict: # feed_dicts are a batch of dicts 112 | feed_dict = super().collate_batch(feed_dicts) 113 | feed_dict['batch_size'] = len(feed_dicts) 114 | predict_dict = self.model.ranker(utils.batch_to_gpu(feed_dict, self.model.device)) # pos+pad+neg+pad 115 | feed_dict['scores'] = predict_dict['prediction'] # [batch(or num_sequence),n_candidate] 116 | pos_mask = torch.arange(0, self.model.train_max_pos_item, device = self.model.device).type_as(feed_dict['pos_num']).unsqueeze(0).expand(feed_dict['batch_size'], self.model.train_max_pos_item).lt(feed_dict['pos_num'].unsqueeze(1)) 117 | neg_mask = torch.arange(0, self.model.train_max_neg_item, device = self.model.device).type_as(feed_dict['neg_num']).unsqueeze(0).expand(feed_dict['batch_size'], self.model.train_max_neg_item).lt(feed_dict['neg_num'].unsqueeze(1)) 118 | all_mask = torch.cat([pos_mask, neg_mask],dim = 1) 119 | feed_dict['padding_mask'] = ~all_mask 120 | feed_dict['scores'] = torch.where(all_mask == 1, feed_dict['scores'],-np.inf * torch.ones_like(feed_dict['scores'])) 121 | _,temp = feed_dict['scores'].sort(dim = 1, descending = True) 122 | _,feed_dict['position'] = temp.sort(dim = 1) 123 | feed_dict['u_v'] = predict_dict['u_v'] # [batch(or num_sequence),ranker_embedding_len] 124 | feed_dict['i_v'] = predict_dict['i_v'] # [batch(or num_sequence),ranker_embedding_len] 125 | 126 | #modeling user history, need all history item vector 127 | ranker = self.model.ranker 128 | if 'LightGCN' in self.model.ranker_name: 129 | all_his_its = ranker.encoder.embedding_dict['item_emb'][feed_dict['history_items'].to(self.model.device),:] 130 | else: 131 | all_his_its = ranker.i_embeddings(feed_dict['history_items'].to(self.model.device)) 132 | feed_dict['his_v']=all_his_its 133 | return feed_dict -------------------------------------------------------------------------------- /src/models/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/THUwangcy/ReChorus/e58d752558634d42be60de561d6f74df748583cc/src/models/__init__.py -------------------------------------------------------------------------------- /src/models/context/AFM.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Zhiyu He 3 | # @Email : hezy22@mails.tsinghua.edu.cn 4 | 5 | """ AFM 6 | Reference: 7 | 'Attentional Factorization Machines: Learning the Weight of Feature Interactions via Attention Networks', Xiao et al, 2017. Arxiv. 8 | Implementation reference: AFM and RecBole 9 | https://github.com/hexiangnan/attentional_factorization_machine 10 | https://github.com/RUCAIBox/RecBole/blob/master/recbole/model/context_aware_recommender/afm.py 11 | """ 12 | 13 | import torch 14 | import torch.nn as nn 15 | import numpy as np 16 | import pandas as pd 17 | 18 | from models.BaseContextModel import ContextCTRModel, ContextModel 19 | from models.context.FM import FMBase 20 | from utils.layers import AttLayer # move here? 21 | 22 | class AFMBase(FMBase): 23 | @staticmethod 24 | def parse_model_args_AFM(parser): 25 | parser.add_argument('--emb_size', type=int, default=64, 26 | help='Size of embedding vectors.') 27 | parser.add_argument('--attention_size', type=int, default=64, 28 | help='Size of attention embedding vectors.') 29 | parser.add_argument('--reg_weight',type=float, default=2.0, 30 | help='Regularization weight for attention layer weights.') 31 | return parser 32 | 33 | def _define_init_afm(self, args, corpus): 34 | self.vec_size = args.emb_size 35 | self.attention_size = args.attention_size 36 | self.reg_weight = args.reg_weight 37 | self._define_params_AFM() 38 | self.apply(self.init_weights) 39 | 40 | def _define_params_AFM(self): 41 | self._define_params_FM() # basic embedding initialization from FM 42 | self.dropout_layer = nn.Dropout(p=self.dropout) 43 | self.attlayer = AttLayer(self.vec_size, self.attention_size) 44 | self.p = torch.nn.Parameter(torch.randn(self.vec_size),requires_grad=True) 45 | 46 | def build_cross(self, feat_emb): 47 | row = [] 48 | col = [] 49 | for i in range(len(self.feature_max)-1): 50 | for j in range(i+1, len(self.feature_max)): 51 | row.append(i) 52 | col.append(j) 53 | p = feat_emb[:,:,row] 54 | q = feat_emb[:,:,col] 55 | return p, q 56 | 57 | def afm_layer(self, infeature): 58 | """Reference: 59 | RecBole - https://github.com/RUCAIBox/RecBole/blob/master/recbole/model/context_aware_recommender/afm.py 60 | """ 61 | p, q = self.build_cross(infeature) 62 | pair_wise_inter = torch.mul(p,q) # batch_size * num_items * num_pairs * emb_dim 63 | 64 | att_signal = self.attlayer(pair_wise_inter).unsqueeze(dim=-1) 65 | att_inter = torch.mul( 66 | att_signal, pair_wise_inter 67 | ) # [batch_size, num_items, num_pairs, emb_dim] 68 | att_pooling = torch.sum(att_inter, dim=-2) # [batch_size, num_items, emb_dim] 69 | att_pooling = self.dropout_layer(att_pooling) # [batch_size, num_items, emb_dim] 70 | 71 | att_pooling = torch.mul(att_pooling, self.p) # [batch_size, num_items, emb_dim] 72 | att_pooling = torch.sum(att_pooling, dim=-1, keepdim=True) # [batch_size, num_items, 1] 73 | return att_pooling 74 | 75 | def forward(self, feed_dict): 76 | fm_vectors, linear_value = self._get_embeddings_FM(feed_dict) 77 | 78 | afm_vectors = self.afm_layer(fm_vectors) 79 | predictions = linear_value + afm_vectors.squeeze(dim=-1) 80 | 81 | return {'prediction':predictions} 82 | 83 | class AFMCTR(ContextCTRModel, AFMBase): # CTR mode 84 | reader, runner = 'ContextReader', 'CTRRunner' 85 | extra_log_args = ['emb_size', 'attention_size', 'loss_n'] 86 | 87 | @staticmethod 88 | def parse_model_args(parser): 89 | parser = AFMBase.parse_model_args_AFM(parser) 90 | return ContextModel.parse_model_args(parser) 91 | 92 | def __init__(self, args, corpus): 93 | ContextCTRModel.__init__(self, args, corpus) 94 | self._define_init_afm(args,corpus) 95 | 96 | def forward(self, feed_dict): 97 | out_dict = AFMBase.forward(self, feed_dict) 98 | out_dict['prediction'] = out_dict['prediction'].view(-1).sigmoid() 99 | out_dict['label'] = feed_dict['label'].view(-1) 100 | return out_dict 101 | 102 | def loss(self, out_dict: dict): 103 | l2_loss = self.reg_weight * torch.norm(self.attlayer.w.weight, p=2) 104 | loss = ContextCTRModel.loss(self, out_dict) 105 | return loss + l2_loss 106 | 107 | class AFMTopK(ContextModel,AFMBase): # Top-k Ranking mode 108 | reader, runner = 'ContextReader', 'BaseRunner' 109 | extra_log_args = ['emb_size', 'attention_size', 'loss_n'] 110 | 111 | @staticmethod 112 | def parse_model_args(parser): 113 | parser = AFMBase.parse_model_args_AFM(parser) 114 | return ContextModel.parse_model_args(parser) 115 | 116 | def __init__(self, args, corpus): 117 | ContextModel.__init__(self, args, corpus) 118 | self._define_init_afm(args,corpus) 119 | 120 | def forward(self, feed_dict): 121 | return AFMBase.forward(self, feed_dict) 122 | 123 | def loss(self, out_dict: dict): 124 | l2_loss = self.reg_weight * torch.norm(self.attlayer.w.weight, p=2) 125 | loss = ContextModel.loss(self, out_dict) 126 | return loss + l2_loss 127 | -------------------------------------------------------------------------------- /src/models/context/AutoInt.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Zhiyu He 3 | # @Email : hezy22@mails.tsinghua.edu.cn 4 | 5 | """AutoInt 6 | Reference: 7 | Weiping Song et al. "AutoInt: Automatic Feature Interaction Learning via Self-Attentive Neural Networks" 8 | in CIKM 2018. 9 | Implementation reference: AutoInt and FuxiCTR 10 | https://github.com/shichence/AutoInt/blob/master/model.py 11 | https://github.com/reczoo/FuxiCTR/blob/main/model_zoo/AutoInt/src/AutoInt.py 12 | """ 13 | 14 | import torch 15 | import torch.nn as nn 16 | import numpy as np 17 | import pandas as pd 18 | 19 | from models.BaseContextModel import ContextCTRModel, ContextModel 20 | from models.context.FM import FMBase 21 | from utils.layers import MultiHeadAttention, MLP_Block 22 | 23 | class AutoIntBase(FMBase): 24 | @staticmethod 25 | def parse_model_args_AutoInt(parser): 26 | parser.add_argument('--emb_size', type=int, default=64, 27 | help='Size of embedding vectors.') 28 | parser.add_argument('--attention_size', type=int, default=32, 29 | help='Size of attention hidden space.') 30 | parser.add_argument('--num_heads', type=int, default=1, 31 | help='Number of attention heads.') 32 | parser.add_argument('--num_layers', type=int, default=1, 33 | help='Number of self-attention layers.') # for attention layer 34 | parser.add_argument('--layers', type=str, default='[64]', 35 | help="Size of each layer.") # for deep layers 36 | return parser 37 | 38 | def _define_init(self, args, corpus): 39 | self.vec_size = args.emb_size 40 | self.layers = eval(args.layers) 41 | 42 | self.num_heads = args.num_heads 43 | self.num_layers = args.num_layers 44 | self.attention_size= args.attention_size 45 | 46 | self._define_params_AutoInt() 47 | self.apply(self.init_weights) 48 | 49 | def _define_params_AutoInt(self): 50 | self._define_params_FM() 51 | # Attention 52 | att_input = self.vec_size 53 | autoint_attentions = [] 54 | residual_embeddings = [] 55 | for _ in range(self.num_layers): 56 | autoint_attentions.append( 57 | MultiHeadAttention(d_model=att_input, n_heads=self.num_heads, kq_same=False, bias=False, 58 | attention_d=self.attention_size)) 59 | residual_embeddings.append(nn.Linear(att_input, self.attention_size)) 60 | att_input = self.attention_size 61 | self.autoint_attentions = nn.ModuleList(autoint_attentions) 62 | self.residual_embeddings = nn.ModuleList(residual_embeddings) 63 | # Deep 64 | pre_size = len(self.feature_max) * self.attention_size 65 | self.deep_layers = MLP_Block(pre_size, self.layers, hidden_activations="ReLU", 66 | dropout_rates=self.dropout, output_dim=1) 67 | 68 | def forward(self, feed_dict): 69 | # FM 70 | autoint_all_embeddings, linear_value = self._get_embeddings_FM(feed_dict) 71 | # Attention + Residual 72 | for autoint_attention, residual_embedding in zip(self.autoint_attentions, self.residual_embeddings): 73 | attention = autoint_attention(autoint_all_embeddings, autoint_all_embeddings, autoint_all_embeddings) 74 | residual = residual_embedding(autoint_all_embeddings) 75 | autoint_all_embeddings = (attention + residual).relu() 76 | # Deep 77 | deep_vectors = autoint_all_embeddings.flatten(start_dim=-2) 78 | deep_vectors = self.deep_layers(deep_vectors) 79 | predictions = linear_value + deep_vectors.squeeze(-1) 80 | return {'prediction':predictions} 81 | 82 | class AutoIntCTR(ContextCTRModel, AutoIntBase): 83 | reader, runner = 'ContextReader', 'CTRRunner' 84 | extra_log_args = ['emb_size','layers','num_layers','num_heads','loss_n'] 85 | 86 | @staticmethod 87 | def parse_model_args(parser): 88 | parser = AutoIntBase.parse_model_args_AutoInt(parser) 89 | return ContextCTRModel.parse_model_args(parser) 90 | 91 | def __init__(self, args, corpus): 92 | ContextCTRModel.__init__(self, args, corpus) 93 | self._define_init(args,corpus) 94 | 95 | def forward(self, feed_dict): 96 | out_dict = AutoIntBase.forward(self, feed_dict) 97 | out_dict['prediction'] = out_dict['prediction'].view(-1).sigmoid() 98 | out_dict['label'] = feed_dict['label'].view(-1) 99 | return out_dict 100 | 101 | class AutoIntTopK(ContextModel,AutoIntBase): 102 | reader, runner = 'ContextReader', 'BaseRunner' 103 | extra_log_args = ['emb_size','layers','num_layers','num_heads','loss_n'] 104 | 105 | @staticmethod 106 | def parse_model_args(parser): 107 | parser = AutoIntBase.parse_model_args_AutoInt(parser) 108 | return ContextModel.parse_model_args(parser) 109 | 110 | def __init__(self, args, corpus): 111 | ContextModel.__init__(self, args, corpus) 112 | self._define_init(args,corpus) 113 | 114 | def forward(self, feed_dict): 115 | return AutoIntBase.forward(self, feed_dict) 116 | -------------------------------------------------------------------------------- /src/models/context/DCN.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Zhiyu He 3 | # @Email : hezy22@mails.tsinghua.edu.cn 4 | 5 | """ DCN 6 | Reference: 7 | 'Deep & Cross Network for Ad Click Predictions', Wang et al, KDD2017. 8 | Implementation reference: RecBole 9 | https://github.com/RUCAIBox/RecBole/blob/master/recbole/model/context_aware_recommender/dcn.py 10 | """ 11 | 12 | import torch 13 | import torch.nn as nn 14 | import torch.nn.functional as fn 15 | import numpy as np 16 | import pandas as pd 17 | 18 | from models.BaseContextModel import ContextCTRModel, ContextModel 19 | from utils.layers import MLP_Block 20 | 21 | class DCNBase(object): 22 | @staticmethod 23 | def parse_model_args_DCN(parser): 24 | parser.add_argument('--emb_size', type=int, default=64, 25 | help='Size of embedding vectors.') 26 | parser.add_argument('--layers', type=str, default='[64]', 27 | help="Size of each deep layer.") 28 | parser.add_argument('--cross_layer_num',type=int,default=6, 29 | help="Number of cross layers.") 30 | parser.add_argument('--reg_weight',type=float, default=2.0, 31 | help="Regularization weight for cross-layer weights. In DCNv2, it is only used for mixed version") 32 | return parser 33 | 34 | def _define_init(self, args, corpus): 35 | self._define_init_params(args,corpus) 36 | self._define_params_DCN() 37 | self.apply(self.init_weights) 38 | 39 | def _define_init_params(self, args, corpus): 40 | self.vec_size = args.emb_size 41 | self.reg_weight = args.reg_weight 42 | self.layers = eval(args.layers) 43 | self.cross_layer_num = args.cross_layer_num 44 | 45 | def _define_params_DCN(self): 46 | # embedding 47 | self.context_embedding = nn.ModuleDict() 48 | for f in self.context_features: 49 | self.context_embedding[f] = nn.Embedding(self.feature_max[f],self.vec_size) if f.endswith('_c') or f.endswith('_id') else\ 50 | nn.Linear(1,self.vec_size,bias=False) 51 | pre_size = len(self.feature_max) * self.vec_size 52 | # cross layers 53 | self.cross_layer_w = nn.ParameterList(nn.Parameter(torch.randn(pre_size),requires_grad=True) 54 | for l in range(self.cross_layer_num)) 55 | self.cross_layer_b = nn.ParameterList(nn.Parameter(torch.tensor([0.01]*pre_size),requires_grad=True) 56 | for l in range(self.cross_layer_num)) 57 | # deep layers 58 | self.deep_layers = MLP_Block(pre_size, self.layers, hidden_activations="ReLU", 59 | batch_norm=True,norm_before_activation=True, 60 | dropout_rates=self.dropout, output_dim=None) 61 | # output layer 62 | self.predict_layer = nn.Linear(len(self.feature_max) * self.vec_size + self.layers[-1], 1) 63 | 64 | def cross_net(self, x_0): 65 | # x_0: batch size * item num * embedding size 66 | ''' 67 | math:: x_{l+1} = x_0 · w_l * x_l^T + b_l + x_l 68 | ''' 69 | x_l = x_0 70 | for layer in range(self.cross_layer_num): 71 | xl_w = torch.tensordot(x_l, self.cross_layer_w[layer], dims=([-1],[0])) 72 | xl_dot = x_0 * xl_w.unsqueeze(2) 73 | x_l = xl_dot + self.cross_layer_b[layer] + x_l 74 | return x_l 75 | 76 | def forward(self, feed_dict): 77 | item_ids = feed_dict['item_id'] 78 | batch_size, item_num = item_ids.shape 79 | # embedding 80 | context_vectors = [self.context_embedding[f](feed_dict[f]) if f.endswith('_c') or f.endswith('_id') 81 | else self.context_embedding[f](feed_dict[f].float().unsqueeze(-1)) 82 | for f in self.context_features] 83 | context_vectors = torch.stack([v if len(v.shape)==3 else v.unsqueeze(dim=-2).repeat(1, item_num, 1) 84 | for v in context_vectors], dim=-2) # batch size * item num * feature num * feature dim 85 | context_emb = context_vectors.flatten(start_dim=-2) 86 | 87 | # cross net 88 | cross_output = self.cross_net(context_emb) 89 | batch_size, item_num, output_emb = cross_output.shape 90 | deep_output = context_emb.view(-1,output_emb) 91 | deep_output = self.deep_layers(deep_output) 92 | deep_output = deep_output.view(batch_size, item_num, self.layers[-1]) 93 | 94 | # output 95 | output = self.predict_layer(torch.cat([cross_output, deep_output],dim=-1)) 96 | predictions = output.squeeze(dim=-1) 97 | return {'prediction':predictions} 98 | 99 | def l2_reg(self, parameters): 100 | """ 101 | Reference: 102 | RecBole - https://github.com/RUCAIBox/RecBole/blob/master/recbole/model/loss.py 103 | RegLoss, L2 regularization on model parameters 104 | """ 105 | reg_loss = None 106 | for W in parameters: 107 | if reg_loss is None: 108 | reg_loss = W.norm(2) 109 | else: 110 | reg_loss = reg_loss + W.norm(2) 111 | return reg_loss 112 | 113 | class DCNCTR(ContextCTRModel, DCNBase): 114 | reader, runner = 'ContextReader', 'CTRRunner' 115 | extra_log_args = ['emb_size','loss_n','cross_layer_num'] 116 | 117 | @staticmethod 118 | def parse_model_args(parser): 119 | parser = DCNBase.parse_model_args_DCN(parser) 120 | return ContextCTRModel.parse_model_args(parser) 121 | 122 | def __init__(self, args, corpus): 123 | ContextCTRModel.__init__(self, args, corpus) 124 | self._define_init(args,corpus) 125 | 126 | def forward(self, feed_dict): 127 | out_dict = DCNBase.forward(self, feed_dict) 128 | out_dict['prediction'] = out_dict['prediction'].view(-1).sigmoid() 129 | out_dict['label'] = feed_dict['label'].view(-1) 130 | return out_dict 131 | 132 | def loss(self, out_dict: dict): 133 | l2_loss = self.reg_weight * DCNBase.l2_reg(self, self.cross_layer_w) 134 | loss = ContextCTRModel.loss(self, out_dict) 135 | return loss + l2_loss 136 | 137 | class DCNTopK(ContextModel, DCNBase): 138 | reader, runner = 'ContextReader', 'BaseRunner' 139 | extra_log_args = ['emb_size','loss_n','cross_layer_num'] 140 | 141 | @staticmethod 142 | def parse_model_args(parser): 143 | parser = DCNBase.parse_model_args_DCN(parser) 144 | return ContextModel.parse_model_args(parser) 145 | 146 | def __init__(self, args, corpus): 147 | ContextModel.__init__(self, args, corpus) 148 | self._define_init(args,corpus) 149 | 150 | def forward(self, feed_dict): 151 | return DCNBase.forward(self, feed_dict) 152 | 153 | def loss(self, out_dict: dict): 154 | l2_loss = self.reg_weight * DCNBase.l2_reg(self, self.cross_layer_w) 155 | loss = ContextModel.loss(self, out_dict) 156 | return loss + l2_loss -------------------------------------------------------------------------------- /src/models/context/DeepFM.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Jiayu Li 3 | # @Email : jy-li20@mails.tsinghua.edu.cn 4 | 5 | """ DeepFM 6 | Reference: 7 | 'DeepFM: A Factorization-Machine based Neural Network for CTR Prediction', Guo et al., IJCAI 2017. 8 | """ 9 | import torch 10 | import torch.nn as nn 11 | import numpy as np 12 | import pandas as pd 13 | 14 | from utils import layers 15 | from models.context.WideDeep import WideDeepCTR, WideDeepTopK 16 | from models.context.WideDeep import WideDeepBase 17 | 18 | class DeepFMBase(WideDeepBase): 19 | def forward(self, feed_dict): 20 | context_vectors, linear_vectors = self._get_embeddings_FM(feed_dict) 21 | # FM 22 | fm_vectors = 0.5 * (context_vectors.sum(dim=-2).pow(2) - context_vectors.pow(2).sum(dim=-2)) 23 | fm_prediction = fm_vectors.sum(dim=-1) + linear_vectors 24 | # Deep 25 | deep_prediction = self.deep_layers(context_vectors.flatten(start_dim=-2)).squeeze(dim=-1) 26 | 27 | predictions = fm_prediction + deep_prediction 28 | return {'prediction':predictions} 29 | 30 | class DeepFMCTR(WideDeepCTR, DeepFMBase): 31 | reader, runner = 'ContextReader', 'CTRRunner' 32 | extra_log_args = ['emb_size','layers','loss_n'] 33 | 34 | def __init__(self, args, corpus): 35 | WideDeepCTR.__init__(self, args, corpus) 36 | 37 | def forward(self, feed_dict): 38 | out_dict = DeepFMBase.forward(self, feed_dict) 39 | out_dict['prediction'] = out_dict['prediction'].view(-1).sigmoid() 40 | out_dict['label'] = feed_dict['label'].view(-1) 41 | return out_dict 42 | 43 | class DeepFMTopK(WideDeepTopK, DeepFMBase): 44 | reader, runner = 'ContextReader','BaseRunner' 45 | extra_log_args = ['emb_size','layers','loss_n'] 46 | 47 | def __init__(self, args, corpus): 48 | WideDeepTopK.__init__(self, args, corpus) 49 | 50 | def forward(self, feed_dict): 51 | return DeepFMBase.forward(self, feed_dict) 52 | -------------------------------------------------------------------------------- /src/models/context/FM.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Jiayu Li 3 | # @Email : jy-li20@mails.tsinghua.edu.cn 4 | 5 | """ FM 6 | Reference: 7 | 'Factorization Machines', Steffen Rendle, 2010 IEEE International conference on data mining. 8 | """ 9 | 10 | import torch 11 | import torch.nn as nn 12 | import numpy as np 13 | import pandas as pd 14 | 15 | from models.BaseContextModel import ContextCTRModel, ContextModel 16 | 17 | class FMBase(object): 18 | @staticmethod 19 | def parse_model_args_FM(parser): 20 | parser.add_argument('--emb_size', type=int, default=64, 21 | help='Size of embedding vectors.') 22 | return parser 23 | 24 | def _define_init_params(self, args,corpus): 25 | self.vec_size = args.emb_size 26 | self._define_params_FM() 27 | self.apply(self.init_weights) 28 | 29 | def _define_init(self, args, corpus): 30 | self._define_init_params(args,corpus) 31 | self._define_params_FM() 32 | self.apply(self.init_weights) 33 | 34 | def _define_params_FM(self): 35 | self.context_embedding = nn.ModuleDict() 36 | self.linear_embedding = nn.ModuleDict() 37 | for f in self.context_features: 38 | self.context_embedding[f] = nn.Embedding(self.feature_max[f],self.vec_size) if f.endswith('_c') or f.endswith('_id') else\ 39 | nn.Linear(1,self.vec_size,bias=False) 40 | self.linear_embedding[f] = nn.Embedding(self.feature_max[f],1) if f.endswith('_c') or f.endswith('_id') else\ 41 | nn.Linear(1,1,bias=False) 42 | self.overall_bias = torch.nn.Parameter(torch.tensor([0.01]), requires_grad=True) 43 | 44 | def _get_embeddings_FM(self, feed_dict): 45 | item_ids = feed_dict['item_id'] 46 | _, item_num = item_ids.shape 47 | 48 | fm_vectors = [self.context_embedding[f](feed_dict[f]) if f.endswith('_c') or f.endswith('_id') 49 | else self.context_embedding[f](feed_dict[f].float().unsqueeze(-1)) for f in self.context_features] 50 | fm_vectors = torch.stack([v if len(v.shape)==3 else v.unsqueeze(dim=-2).repeat(1, item_num, 1) 51 | for v in fm_vectors], dim=-2) # batch size * item num * feature num * feature dim: 84,100,2,64 52 | linear_value = [self.linear_embedding[f](feed_dict[f]) if f.endswith('_c') or f.endswith('_id') 53 | else self.linear_embedding[f](feed_dict[f].float().unsqueeze(-1)) for f in self.context_features] 54 | linear_value = torch.cat([v if len(v.shape)==3 else v.unsqueeze(dim=-2).repeat(1, item_num, 1) 55 | for v in linear_value],dim=-1) # batch size * item num * feature num 56 | linear_value = self.overall_bias + linear_value.sum(dim=-1) 57 | return fm_vectors, linear_value 58 | 59 | def forward(self, feed_dict): 60 | fm_vectors, linear_value = self._get_embeddings_FM(feed_dict) 61 | fm_vectors = 0.5 * (fm_vectors.sum(dim=-2).pow(2) - fm_vectors.pow(2).sum(dim=-2)) 62 | predictions = linear_value + fm_vectors.sum(dim=-1) 63 | return {'prediction':predictions} 64 | 65 | class FMCTR(ContextCTRModel, FMBase): 66 | reader, runner = 'ContextReader', 'CTRRunner' 67 | extra_log_args = ['emb_size','loss_n'] 68 | 69 | @staticmethod 70 | def parse_model_args(parser): 71 | parser = FMBase.parse_model_args_FM(parser) 72 | return ContextCTRModel.parse_model_args(parser) 73 | 74 | def __init__(self, args, corpus): 75 | ContextCTRModel.__init__(self, args, corpus) 76 | self._define_init(args,corpus) 77 | 78 | def forward(self, feed_dict): 79 | out_dict = FMBase.forward(self, feed_dict) 80 | out_dict['prediction'] = out_dict['prediction'].view(-1).sigmoid() 81 | out_dict['label'] = feed_dict['label'].view(-1) 82 | return out_dict 83 | 84 | class FMTopK(ContextModel,FMBase): 85 | reader, runner = 'ContextReader', 'BaseRunner' 86 | extra_log_args = ['emb_size','loss_n'] 87 | 88 | @staticmethod 89 | def parse_model_args(parser): 90 | parser = FMBase.parse_model_args_FM(parser) 91 | return ContextModel.parse_model_args(parser) 92 | 93 | def __init__(self, args, corpus): 94 | ContextModel.__init__(self, args, corpus) 95 | self._define_init(args,corpus) 96 | 97 | def forward(self, feed_dict): 98 | return FMBase.forward(self, feed_dict) -------------------------------------------------------------------------------- /src/models/context/WideDeep.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Jiayu Li 3 | # @Email : jy-li20@mails.tsinghua.edu.cn 4 | 5 | """ WideDeep 6 | Reference: 7 | Wide {\&} Deep Learning for Recommender Systems, Cheng et al. 2016. The 1st workshop on deep learning for recommender systems. 8 | """ 9 | 10 | import torch 11 | import torch.nn as nn 12 | import numpy as np 13 | import pandas as pd 14 | 15 | from models.BaseContextModel import ContextModel, ContextCTRModel 16 | from models.context.FM import FMBase 17 | from utils.layers import MLP_Block 18 | 19 | class WideDeepBase(FMBase): 20 | @staticmethod 21 | def parse_model_args_WD(parser): 22 | parser.add_argument('--emb_size', type=int, default=64, 23 | help='Size of embedding vectors.') 24 | parser.add_argument('--layers', type=str, default='[64]', 25 | help="Size of each layer.") 26 | return parser 27 | 28 | def _define_init(self, args, corpus): 29 | self._define_init_params(args,corpus) 30 | self.layers = eval(args.layers) 31 | self._define_params_WD() 32 | self.apply(self.init_weights) 33 | 34 | def _define_params_WD(self): 35 | self._define_params_FM() 36 | pre_size = len(self.context_features) * self.vec_size 37 | # deep layers 38 | self.deep_layers = MLP_Block(pre_size, self.layers, hidden_activations="ReLU", 39 | batch_norm=False, dropout_rates=self.dropout, output_dim=1) 40 | 41 | def forward(self, feed_dict): 42 | deep_vectors, wide_prediction = self._get_embeddings_FM(feed_dict) 43 | deep_vector = deep_vectors.flatten(start_dim=-2) 44 | deep_prediction = self.deep_layers(deep_vector).squeeze(dim=-1) 45 | predictions = deep_prediction + wide_prediction 46 | return {'prediction':predictions} 47 | 48 | class WideDeepCTR(ContextCTRModel, WideDeepBase): 49 | reader, runner = 'ContextReader', 'CTRRunner' 50 | extra_log_args = ['emb_size','layers','loss_n'] 51 | @staticmethod 52 | def parse_model_args(parser): 53 | parser = WideDeepBase.parse_model_args_WD(parser) 54 | return ContextModel.parse_model_args(parser) 55 | 56 | def __init__(self, args, corpus): 57 | ContextCTRModel.__init__(self, args, corpus) 58 | self._define_init(args,corpus) 59 | 60 | def forward(self, feed_dict): 61 | out_dict = WideDeepBase.forward(self, feed_dict) 62 | out_dict['prediction'] = out_dict['prediction'].view(-1).sigmoid() 63 | out_dict['label'] = feed_dict['label'].view(-1) 64 | return out_dict 65 | 66 | class WideDeepTopK(ContextModel,WideDeepBase): 67 | reader, runner = 'ContextReader','BaseRunner' 68 | extra_log_args = ['emb_size','layers','loss_n'] 69 | 70 | @staticmethod 71 | def parse_model_args(parser): 72 | parser = WideDeepBase.parse_model_args_WD(parser) 73 | return ContextModel.parse_model_args(parser) 74 | 75 | def __init__(self, args, corpus): 76 | ContextModel.__init__(self, args, corpus) 77 | self._define_init(args, corpus) 78 | 79 | def forward(self, feed_dict): 80 | return WideDeepBase.forward(self, feed_dict) 81 | -------------------------------------------------------------------------------- /src/models/context/__init__.py: -------------------------------------------------------------------------------- 1 | from os.path import dirname, basename, isfile, join 2 | import glob 3 | 4 | modules = glob.glob(join(dirname(__file__), "*.py")) 5 | __all__ = [ 6 | basename(f)[:-3] for f in modules if isfile(f) and not f.endswith('__init__.py') 7 | ] 8 | -------------------------------------------------------------------------------- /src/models/context_seq/__init__.py: -------------------------------------------------------------------------------- 1 | from os.path import dirname, basename, isfile, join 2 | import glob 3 | 4 | modules = glob.glob(join(dirname(__file__), "*.py")) 5 | __all__ = [ 6 | basename(f)[:-3] for f in modules if isfile(f) and not f.endswith('__init__.py') 7 | ] 8 | -------------------------------------------------------------------------------- /src/models/developing/CLRec.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import torch 4 | import torch.nn as nn 5 | import torch.nn.functional as F 6 | import numpy as np 7 | 8 | from models.BaseModel import SequentialModel 9 | from utils import layers 10 | 11 | 12 | class CLRec(SequentialModel): 13 | reader = 'SeqReader' 14 | runner = 'BaseRunner' 15 | extra_log_args = ['batch_size', 'temp'] 16 | 17 | @staticmethod 18 | def parse_model_args(parser): 19 | parser.add_argument('--emb_size', type=int, default=64, 20 | help='Size of embedding vectors.') 21 | parser.add_argument('--temp', type=float, default=0.2, 22 | help='Temperature in contrastive loss.') 23 | return SequentialModel.parse_model_args(parser) 24 | 25 | def __init__(self, args, corpus): 26 | super().__init__(args, corpus) 27 | self.emb_size = args.emb_size 28 | self.temp = args.temp 29 | self.max_his = args.history_max 30 | self._define_params() 31 | self.apply(self.init_weights) 32 | 33 | def _define_params(self): 34 | self.i_embeddings = nn.Embedding(self.item_num, self.emb_size) 35 | self.encoder = BERT4RecEncoder(self.emb_size, self.max_his, num_layers=2, num_heads=2) 36 | self.contra_loss = ContraLoss(temperature=self.temp) 37 | 38 | def forward(self, feed_dict): 39 | self.check_list = [] 40 | i_ids = feed_dict['item_id'] # bsz, n_candidate 41 | history = feed_dict['history_items'] # bsz, history_max 42 | lengths = feed_dict['lengths'] # bsz 43 | 44 | his_vectors = self.i_embeddings(history) 45 | his_vector = self.encoder(his_vectors, lengths) 46 | i_vectors = self.i_embeddings(i_ids) 47 | # his_vector = F.normalize(his_vector, dim=-1) 48 | # i_vectors = F.normalize(i_vectors, dim=-1) 49 | prediction = (his_vector[:, None, :] * i_vectors).sum(-1) 50 | out_dict = {'prediction': prediction} 51 | 52 | if feed_dict['phase'] == 'train': 53 | target_vector = i_vectors[:, 0, :] 54 | features = torch.stack([his_vector, target_vector], dim=1) # bsz, 2, emb 55 | features = F.normalize(features, dim=-1) 56 | out_dict['features'] = features # bsz, 2, emb 57 | 58 | return out_dict 59 | 60 | def loss(self, out_dict): 61 | return self.contra_loss(out_dict['features']) 62 | 63 | class Dataset(SequentialModel.Dataset): 64 | # No need to sample negative items 65 | def actions_before_epoch(self): 66 | self.data['neg_items'] = [[] for _ in range(len(self))] 67 | 68 | 69 | """ Contrastive Loss """ 70 | class ContraLoss(nn.Module): 71 | def __init__(self, temperature=0.2): 72 | super(ContraLoss, self).__init__() 73 | self.temperature = temperature 74 | 75 | def forward(self, features, mask=None): 76 | """ 77 | Args: 78 | features: hidden vector of shape [bsz, n_views, ...]. 79 | mask: contrastive mask of shape [bsz, bsz], mask_{i,j}=1 if sequence j 80 | has the same target item as sequence i. Can be asymmetric. 81 | Returns: 82 | A loss scalar. 83 | """ 84 | if len(features.shape) < 3: 85 | raise ValueError('`features` needs to be [bsz, n_views, ...],' 86 | 'at least 3 dimensions are required') 87 | if len(features.shape) > 3: 88 | features = features.view(features.shape[0], features.shape[1], -1) 89 | 90 | batch_size, device = features.shape[0], features.device 91 | if mask is None: 92 | mask = torch.eye(batch_size, dtype=torch.float32).to(device) 93 | 94 | # compute logits 95 | dot_contrast = torch.matmul(features[:, 0], features[:, 1].transpose(0, 1)) / self.temperature 96 | # for numerical stability 97 | logits_max, _ = torch.max(dot_contrast, dim=1, keepdim=True) 98 | logits = dot_contrast - logits_max.detach() # bsz, bsz 99 | 100 | # compute log_prob 101 | exp_logits = torch.exp(logits) 102 | log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True) + 1e-10) 103 | 104 | # compute mean of log-likelihood over positive 105 | mean_log_prob_pos = (mask * log_prob).sum(1) 106 | 107 | return -mean_log_prob_pos.mean() 108 | 109 | 110 | """ Encoder Layer """ 111 | class BERT4RecEncoder(nn.Module): 112 | def __init__(self, emb_size, max_his, num_layers=2, num_heads=2): 113 | super().__init__() 114 | self.p_embeddings = nn.Embedding(max_his + 1, emb_size) 115 | self.transformer_block = nn.ModuleList([ 116 | layers.TransformerLayer(d_model=emb_size, d_ff=emb_size, n_heads=num_heads) 117 | for _ in range(num_layers) 118 | ]) 119 | 120 | def forward(self, seq, lengths): 121 | batch_size, seq_len = seq.size(0), seq.size(1) 122 | len_range = torch.from_numpy(np.arange(seq_len)).to(seq.device) 123 | valid_mask = len_range[None, :] < lengths[:, None] 124 | 125 | # Position embedding 126 | position = len_range[None, :] * valid_mask.long() 127 | pos_vectors = self.p_embeddings(position) 128 | seq = seq + pos_vectors 129 | 130 | # Self-attention 131 | attn_mask = valid_mask.view(batch_size, 1, 1, seq_len) 132 | for block in self.transformer_block: 133 | seq = block(seq, attn_mask) 134 | seq = seq * valid_mask[:, :, None].float() 135 | 136 | his_vector = seq[torch.arange(batch_size), lengths - 1] 137 | return his_vector 138 | -------------------------------------------------------------------------------- /src/models/developing/FourierTA.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import torch 4 | import torch.nn as nn 5 | import numpy as np 6 | 7 | from utils import layers 8 | from models.BaseModel import SequentialModel 9 | from helpers.KDAReader import KDAReader 10 | 11 | 12 | class FourierTA(SequentialModel): 13 | reader = 'SeqReader' 14 | runner = 'BaseRunner' 15 | extra_log_args = ['t_scalar'] 16 | 17 | @staticmethod 18 | def parse_model_args(parser): 19 | parser.add_argument('--emb_size', type=int, default=64, 20 | help='Size of embedding vectors.') 21 | parser.add_argument('--t_scalar', type=int, default=60, 22 | help='Time interval scalar.') 23 | return SequentialModel.parse_model_args(parser) 24 | 25 | def __init__(self, args, corpus): 26 | super().__init__(args, corpus) 27 | self.freq_dim = args.emb_size 28 | self.emb_size = args.emb_size 29 | self.t_scalar = args.t_scalar 30 | self._define_params() 31 | self.apply(self.init_weights) 32 | 33 | def _define_params(self): 34 | self.user_embeddings = nn.Embedding(self.user_num, self.emb_size) 35 | self.item_embeddings = nn.Embedding(self.item_num, self.emb_size) 36 | 37 | self.fourier_attn = FourierTemporalAttention(self.emb_size, self.freq_dim, self.device) 38 | self.W1 = nn.Linear(self.emb_size, self.emb_size) 39 | self.W2 = nn.Linear(self.emb_size, self.emb_size) 40 | self.dropout_layer = nn.Dropout(self.dropout) 41 | self.layer_norm = nn.LayerNorm(self.emb_size) 42 | self.item_bias = nn.Embedding(self.item_num, 1) 43 | 44 | def forward(self, feed_dict): 45 | self.check_list = [] 46 | u_ids = feed_dict['user_id'] # B 47 | i_ids = feed_dict['item_id'] # B * -1 48 | history = feed_dict['history_items'] # B * H 49 | delta_t_n = feed_dict['history_delta_t'].float() # B * H 50 | batch_size, seq_len = history.shape 51 | 52 | u_vectors = self.user_embeddings(u_ids) 53 | i_vectors = self.item_embeddings(i_ids) 54 | his_vectors = self.item_embeddings(history) # B * H * V 55 | 56 | valid_mask = (history > 0).view(batch_size, 1, seq_len) 57 | context = self.fourier_attn(his_vectors, delta_t_n, i_vectors, valid_mask) # B * -1 * V 58 | 59 | residual = context 60 | # feed forward 61 | context = self.W1(context) 62 | context = self.W2(context.relu()) 63 | # dropout, residual and layer_norm 64 | context = self.dropout_layer(context) 65 | context = self.layer_norm(residual + context) 66 | # context = self.layer_norm(context) 67 | 68 | i_bias = self.item_bias(i_ids).squeeze(-1) 69 | prediction = ((u_vectors[:, None, :] + context) * i_vectors).sum(dim=-1) 70 | prediction = prediction + i_bias 71 | out_dict = {'prediction': prediction} 72 | return out_dict 73 | 74 | class Dataset(SequentialModel.Dataset): 75 | def _get_feed_dict(self, index): 76 | feed_dict = super()._get_feed_dict(index) 77 | delta_t = self.data['time'][index] - feed_dict['history_times'] 78 | feed_dict['history_delta_t'] = KDAReader.norm_time(delta_t, self.model.t_scalar) 79 | return feed_dict 80 | 81 | 82 | class FourierTemporalAttention(nn.Module): 83 | def __init__(self, emb_size: int, freq_dim: int, device): 84 | super().__init__() 85 | self.d = emb_size 86 | self.d_f = freq_dim 87 | 88 | self.freq_real = nn.Parameter(torch.zeros(self.d_f)) 89 | self.freq_imag = nn.Parameter(torch.zeros(self.d_f)) 90 | self.A = nn.Linear(self.d, 10) 91 | self.A_out = nn.Linear(10, 1, bias=False) 92 | 93 | nn.init.normal_(self.freq_real.data, mean=0.0, std=0.01) 94 | nn.init.normal_(self.freq_imag.data, mean=0.0, std=0.01) 95 | freq = np.linspace(0, 1, self.d_f) / 2. 96 | self.freqs = torch.from_numpy(np.concatenate((freq, -freq))).to(device).float() 97 | 98 | def idft_decay(self, delta_t): 99 | # create conjugate symmetric to ensure real number output 100 | x_real = torch.cat([self.freq_real, self.freq_real], dim=-1) 101 | x_imag = torch.cat([self.freq_imag, -self.freq_imag], dim=-1) 102 | w = 2. * np.pi * self.freqs * delta_t.unsqueeze(-1) # B * H * n_freq 103 | real_part = w.cos() * x_real[None, None, :] # B * H * n_freq 104 | imag_part = w.sin() * x_imag[None, None, :] 105 | decay = (real_part - imag_part).mean(dim=-1) / 2. # B * H 106 | return decay.clamp(0, 1).float() 107 | 108 | def forward(self, seq, delta_t_n, target, valid_mask): 109 | query_vector = seq[:, None, :, :] * target[:, :, None, :] 110 | attention = self.A_out(self.A(query_vector).tanh()).squeeze(-1) # B * -1 * H 111 | # attention = torch.matmul(target, seq.transpose(-2, -1)) / self.d ** 0.5 # B * -1 * H 112 | # shift masked softmax 113 | attention = attention - attention.max() 114 | attention = attention.masked_fill(valid_mask==0, -np.inf).softmax(dim=-1) 115 | # temporal evolution 116 | decay = self.idft_decay(delta_t_n).unsqueeze(1).masked_fill(valid_mask==0, 0.) # B * 1 * H 117 | attention = attention * decay 118 | # attentional aggregation of history items 119 | context = torch.matmul(attention, seq) # B * -1 * V 120 | return context 121 | -------------------------------------------------------------------------------- /src/models/developing/SRGNN.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import torch 4 | from torch import nn 5 | from torch.nn import Parameter 6 | from torch.nn import functional as F 7 | import numpy as np 8 | 9 | from models.BaseModel import SequentialModel 10 | 11 | 12 | class SRGNN(SequentialModel): 13 | reader = 'SeqReader' 14 | runner = 'BaseRunner' 15 | extra_log_args = ['num_layers'] 16 | 17 | @staticmethod 18 | def parse_model_args(parser): 19 | parser.add_argument('--emb_size', type=int, default=64, 20 | help='Size of embedding vectors.') 21 | parser.add_argument('--num_layers', type=int, default=1, 22 | help='Number of self-attention layers.') 23 | return SequentialModel.parse_model_args(parser) 24 | 25 | def __init__(self, args, corpus): 26 | super().__init__(args, corpus) 27 | self.emb_size = args.emb_size 28 | self.num_layers = args.num_layers 29 | self._define_params() 30 | std = 1.0 / np.sqrt(self.emb_size) 31 | for weight in self.parameters(): 32 | weight.data.uniform_(-std, std) 33 | 34 | def _define_params(self): 35 | self.i_embeddings = nn.Embedding(self.item_num, self.emb_size, padding_idx=0) 36 | self.linear1 = nn.Linear(self.emb_size, self.emb_size, bias=True) 37 | self.linear2 = nn.Linear(self.emb_size, self.emb_size, bias=True) 38 | self.linear3 = nn.Linear(self.emb_size, 1, bias=False) 39 | self.linear_transform = nn.Linear(self.emb_size * 2, self.emb_size, bias=True) 40 | self.gnn = GNN(self.emb_size, self.num_layers) 41 | 42 | def _get_slice(self, item_seq): 43 | items, n_node, A, alias_inputs = [], [], [], [] 44 | max_n_node = item_seq.size(1) 45 | item_seq = item_seq.cpu().numpy() 46 | for u_input in item_seq: 47 | node = np.unique(u_input) 48 | items.append(node.tolist() + [0] * (max_n_node - len(node))) 49 | u_A = np.zeros((max_n_node, max_n_node)) 50 | 51 | for i in np.arange(len(u_input) - 1): 52 | if u_input[i + 1] == 0: 53 | break 54 | u = np.where(node == u_input[i])[0][0] 55 | v = np.where(node == u_input[i + 1])[0][0] 56 | u_A[u][v] = 1 57 | 58 | u_sum_in = np.sum(u_A, 0) 59 | u_sum_in[np.where(u_sum_in == 0)] = 1 60 | u_A_in = np.divide(u_A, u_sum_in) 61 | u_sum_out = np.sum(u_A, 1) 62 | u_sum_out[np.where(u_sum_out == 0)] = 1 63 | u_A_out = np.divide(u_A.transpose(), u_sum_out) 64 | u_A = np.concatenate([u_A_in, u_A_out]).transpose() 65 | A.append(u_A) 66 | 67 | alias_inputs.append([np.where(node == i)[0][0] for i in u_input]) 68 | # The relative coordinates of the item node, shape of [batch_size, max_session_len] 69 | alias_inputs = torch.LongTensor(alias_inputs).to(self.device) 70 | # The connecting matrix, shape of [batch_size, max_session_len, 2 * max_session_len] 71 | A = torch.FloatTensor(A).to(self.device) 72 | # The unique item nodes, shape of [batch_size, max_session_len] 73 | items = torch.LongTensor(items).to(self.device) 74 | 75 | return alias_inputs, A, items 76 | 77 | def forward(self, feed_dict): 78 | self.check_list = [] 79 | i_ids = feed_dict['item_id'] # [batch_size, -1] 80 | history = feed_dict['history_items'] # [batch_size, history_max] 81 | lengths = feed_dict['lengths'] # [batch_size] 82 | batch_size, seq_len = history.shape 83 | 84 | valid_his = (history > 0).long() 85 | alias_inputs, A, items = self._get_slice(history) 86 | hidden = self.i_embeddings(items) 87 | hidden = self.gnn(A, hidden) 88 | alias_inputs = alias_inputs.unsqueeze(-1).expand(-1, -1, self.emb_size) 89 | seq_hidden = torch.gather(hidden, dim=1, index=alias_inputs) 90 | 91 | # fetch the last hidden state of last timestamp 92 | ht = seq_hidden[torch.arange(batch_size), lengths - 1] 93 | alpha = self.linear3((self.linear1(ht)[:, None, :] + self.linear2(seq_hidden)).sigmoid()) 94 | a = torch.sum(alpha * seq_hidden * valid_his[:, :, None].float(), 1) 95 | his_vector = self.linear_transform(torch.cat([a, ht], dim=1)) 96 | 97 | i_vectors = self.i_embeddings(i_ids) 98 | prediction = (his_vector[:, None, :] * i_vectors).sum(-1) 99 | return {'prediction': prediction.view(batch_size, -1)} 100 | 101 | 102 | class GNN(nn.Module): 103 | """ 104 | Graph neural networks are well-suited for session-based recommendation, 105 | because it can automatically extract features of session graphs with considerations of rich node connections. 106 | """ 107 | 108 | def __init__(self, embedding_size, step=1): 109 | super(GNN, self).__init__() 110 | self.step = step 111 | self.embedding_size = embedding_size 112 | self.input_size = embedding_size * 2 113 | self.gate_size = embedding_size * 3 114 | self.w_ih = Parameter(torch.Tensor(self.gate_size, self.input_size)) 115 | self.w_hh = Parameter(torch.Tensor(self.gate_size, self.embedding_size)) 116 | self.b_ih = Parameter(torch.Tensor(self.gate_size)) 117 | self.b_hh = Parameter(torch.Tensor(self.gate_size)) 118 | self.b_iah = Parameter(torch.Tensor(self.embedding_size)) 119 | self.b_ioh = Parameter(torch.Tensor(self.embedding_size)) 120 | 121 | self.linear_edge_in = nn.Linear(self.embedding_size, self.embedding_size, bias=True) 122 | self.linear_edge_out = nn.Linear(self.embedding_size, self.embedding_size, bias=True) 123 | 124 | def gnn_cell(self, A, hidden): 125 | """Obtain latent vectors of nodes via graph neural networks. 126 | Args: 127 | A(torch.FloatTensor):The connection matrix,shape of [batch_size, max_session_len, 2 * max_session_len] 128 | hidden(torch.FloatTensor):The item node embedding matrix, shape of 129 | [batch_size, max_session_len, embedding_size] 130 | Returns: 131 | torch.FloatTensor:Latent vectors of nodes,shape of [batch_size, max_session_len, embedding_size] 132 | """ 133 | input_in = torch.matmul(A[:, :, :A.size(1)], self.linear_edge_in(hidden)) + self.b_iah 134 | input_out = torch.matmul(A[:, :, A.size(1): 2 * A.size(1)], self.linear_edge_out(hidden)) + self.b_ioh 135 | # [batch_size, max_session_len, embedding_size * 2] 136 | inputs = torch.cat([input_in, input_out], 2) 137 | 138 | # gi.size equals to gh.size, shape of [batch_size, max_session_len, embdding_size * 3] 139 | gi = F.linear(inputs, self.w_ih, self.b_ih) 140 | gh = F.linear(hidden, self.w_hh, self.b_hh) 141 | # (batch_size, max_session_len, embedding_size) 142 | i_r, i_i, i_n = gi.chunk(3, 2) 143 | h_r, h_i, h_n = gh.chunk(3, 2) 144 | resetgate = torch.sigmoid(i_r + h_r) 145 | inputgate = torch.sigmoid(i_i + h_i) 146 | newgate = torch.tanh(i_n + resetgate * h_n) 147 | hy = (1 - inputgate) * hidden + inputgate * newgate 148 | return hy 149 | 150 | def forward(self, A, hidden): 151 | for i in range(self.step): 152 | hidden = self.gnn_cell(A, hidden) 153 | return hidden 154 | -------------------------------------------------------------------------------- /src/models/developing/__init__.py: -------------------------------------------------------------------------------- 1 | from os.path import dirname, basename, isfile, join 2 | import glob 3 | 4 | modules = glob.glob(join(dirname(__file__), "*.py")) 5 | __all__ = [ 6 | basename(f)[:-3] for f in modules if isfile(f) and not f.endswith('__init__.py') 7 | ] 8 | -------------------------------------------------------------------------------- /src/models/general/BPRMF.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Chenyang Wang 3 | # @Email : THUwangcy@gmail.com 4 | 5 | """ BPRMF 6 | Reference: 7 | "Bayesian personalized ranking from implicit feedback" 8 | Rendle et al., UAI'2009. 9 | CMD example: 10 | python main.py --model_name BPRMF --emb_size 64 --lr 1e-3 --l2 1e-6 --dataset 'Grocery_and_Gourmet_Food' 11 | """ 12 | 13 | import torch.nn as nn 14 | 15 | from models.BaseModel import GeneralModel 16 | from models.BaseImpressionModel import ImpressionModel 17 | 18 | class BPRMFBase(object): 19 | @staticmethod 20 | def parse_model_args(parser): 21 | parser.add_argument('--emb_size', type=int, default=64, 22 | help='Size of embedding vectors.') 23 | return parser 24 | 25 | def _base_init(self, args, corpus): 26 | self.emb_size = args.emb_size 27 | self._base_define_params() 28 | self.apply(self.init_weights) 29 | 30 | def _base_define_params(self): 31 | self.u_embeddings = nn.Embedding(self.user_num, self.emb_size) 32 | self.i_embeddings = nn.Embedding(self.item_num, self.emb_size) 33 | 34 | def forward(self, feed_dict): 35 | self.check_list = [] 36 | u_ids = feed_dict['user_id'] # [batch_size] 37 | i_ids = feed_dict['item_id'] # [batch_size, -1] 38 | 39 | cf_u_vectors = self.u_embeddings(u_ids) 40 | cf_i_vectors = self.i_embeddings(i_ids) 41 | 42 | prediction = (cf_u_vectors[:, None, :] * cf_i_vectors).sum(dim=-1) # [batch_size, -1] 43 | u_v = cf_u_vectors.repeat(1,i_ids.shape[1]).view(i_ids.shape[0],i_ids.shape[1],-1) 44 | i_v = cf_i_vectors 45 | return {'prediction': prediction.view(feed_dict['batch_size'], -1), 'u_v': u_v, 'i_v':i_v} 46 | 47 | class BPRMF(GeneralModel, BPRMFBase): 48 | reader = 'BaseReader' 49 | runner = 'BaseRunner' 50 | extra_log_args = ['emb_size', 'batch_size'] 51 | 52 | @staticmethod 53 | def parse_model_args(parser): 54 | parser = BPRMFBase.parse_model_args(parser) 55 | return GeneralModel.parse_model_args(parser) 56 | 57 | def __init__(self, args, corpus): 58 | GeneralModel.__init__(self, args, corpus) 59 | self._base_init(args, corpus) 60 | 61 | def forward(self, feed_dict): 62 | out_dict = BPRMFBase.forward(self, feed_dict) 63 | return {'prediction': out_dict['prediction']} 64 | 65 | class BPRMFImpression(ImpressionModel, BPRMFBase): 66 | reader = 'ImpressionReader' 67 | runner = 'ImpressionRunner' 68 | extra_log_args = ['emb_size', 'batch_size'] 69 | 70 | @staticmethod 71 | def parse_model_args(parser): 72 | parser = BPRMFBase.parse_model_args(parser) 73 | return ImpressionModel.parse_model_args(parser) 74 | 75 | def __init__(self, args, corpus): 76 | ImpressionModel.__init__(self, args, corpus) 77 | self._base_init(args, corpus) 78 | 79 | def forward(self, feed_dict): 80 | return BPRMFBase.forward(self, feed_dict) -------------------------------------------------------------------------------- /src/models/general/BUIR.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Chenyang Wang 3 | # @Email : THUwangcy@gmail.com 4 | 5 | """ BUIR 6 | Reference: 7 | "Bootstrapping User and Item Representations for One-Class Collaborative Filtering" 8 | Lee et al., SIGIR'2021. 9 | CMD example: 10 | python main.py --model_name BUIR --emb_size 64 --lr 1e-3 --l2 1e-6 --dataset 'Grocery_and_Gourmet_Food' 11 | """ 12 | 13 | import numpy as np 14 | import torch.nn as nn 15 | import torch.nn.functional as F 16 | 17 | from models.BaseModel import GeneralModel 18 | 19 | 20 | class BUIR(GeneralModel): 21 | reader = 'BaseReader' 22 | runner = 'BUIRRunner' 23 | extra_log_args = ['emb_size', 'momentum'] 24 | 25 | @staticmethod 26 | def parse_model_args(parser): 27 | parser.add_argument('--emb_size', type=int, default=64, 28 | help='Size of embedding vectors.') 29 | parser.add_argument('--momentum', type=float, default=0.995, 30 | help='Momentum update.') 31 | return GeneralModel.parse_model_args(parser) 32 | 33 | @staticmethod 34 | def init_weights(m): 35 | if 'Linear' in str(type(m)): 36 | nn.init.xavier_normal_(m.weight.data) 37 | if m.bias is not None: 38 | nn.init.normal_(m.bias.data) 39 | elif 'Embedding' in str(type(m)): 40 | nn.init.xavier_normal_(m.weight.data) 41 | 42 | def __init__(self, args, corpus): 43 | super().__init__(args, corpus) 44 | self.emb_size = args.emb_size 45 | self.momentum = args.momentum 46 | self._define_params() 47 | self.apply(self.init_weights) 48 | 49 | for param_o, param_t in zip(self.user_online.parameters(), self.user_target.parameters()): 50 | param_t.data.copy_(param_o.data) 51 | param_t.requires_grad = False 52 | for param_o, param_t in zip(self.item_online.parameters(), self.item_target.parameters()): 53 | param_t.data.copy_(param_o.data) 54 | param_t.requires_grad = False 55 | 56 | 57 | def _define_params(self): 58 | self.user_online = nn.Embedding(self.user_num, self.emb_size) 59 | self.user_target = nn.Embedding(self.user_num, self.emb_size) 60 | self.item_online = nn.Embedding(self.item_num, self.emb_size) 61 | self.item_target = nn.Embedding(self.item_num, self.emb_size) 62 | self.predictor = nn.Linear(self.emb_size, self.emb_size) 63 | self.bn = nn.BatchNorm1d(self.emb_size, eps=0, affine=False, track_running_stats=False) 64 | 65 | # will be called by BUIRRunner 66 | def _update_target(self): 67 | for param_o, param_t in zip(self.user_online.parameters(), self.user_target.parameters()): 68 | param_t.data = param_t.data * self.momentum + param_o.data * (1. - self.momentum) 69 | 70 | for param_o, param_t in zip(self.item_online.parameters(), self.item_target.parameters()): 71 | param_t.data = param_t.data * self.momentum + param_o.data * (1. - self.momentum) 72 | 73 | def forward(self, feed_dict): 74 | self.check_list = [] 75 | user, items = feed_dict['user_id'], feed_dict['item_id'] 76 | 77 | # prediction = (self.item_online(items) * self.user_online(user)[:, None, :]).sum(-1) 78 | prediction = (self.predictor(self.item_online(items)) * self.user_online(user)[:, None, :]).sum(dim=-1) + \ 79 | (self.predictor(self.user_online(user))[:, None, :] * self.item_online(items)).sum(dim=-1) 80 | out_dict = {'prediction': prediction} 81 | 82 | if feed_dict['phase'] == 'train': 83 | u_online = self.user_online(user) 84 | u_online = self.predictor(u_online) 85 | u_target = self.user_target(user) 86 | i_online = self.item_online(items).squeeze(1) 87 | i_online = self.predictor(i_online) 88 | i_target = self.item_target(items).squeeze(1) 89 | out_dict.update({ 90 | 'u_online': u_online, 91 | 'u_target': u_target, 92 | 'i_online': i_online, 93 | 'i_target': i_target 94 | }) 95 | return out_dict 96 | 97 | def loss(self, output): 98 | u_online, u_target = output['u_online'], output['u_target'] 99 | i_online, i_target = output['i_online'], output['i_target'] 100 | 101 | u_online = F.normalize(u_online, dim=-1) 102 | u_target = F.normalize(u_target, dim=-1) 103 | i_online = F.normalize(i_online, dim=-1) 104 | i_target = F.normalize(i_target, dim=-1) 105 | 106 | # Euclidean distance between normalized vectors can be replaced with their negative inner product 107 | loss_ui = 2 - 2 * (u_online * i_target.detach()).sum(dim=-1) 108 | loss_iu = 2 - 2 * (i_online * u_target.detach()).sum(dim=-1) 109 | 110 | return (loss_ui + loss_iu).mean() 111 | 112 | class Dataset(GeneralModel.Dataset): 113 | # No need to sample negative items 114 | def actions_before_epoch(self): 115 | self.data['neg_items'] = [[] for _ in range(len(self))] 116 | -------------------------------------------------------------------------------- /src/models/general/CFKG.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Chenyang Wang 3 | # @Email : THUwangcy@gmail.com 4 | 5 | """ CFKG 6 | Reference: 7 | "Learning over Knowledge-Base Embeddings for Recommendation" 8 | Yongfeng Zhang et al., SIGIR'2018. 9 | Note: 10 | In the built-in dataset, we have four kinds of relations: buy, category, complement, substitute, where 'buy' is 11 | a special relation indexed by 0. And there are three kinds of nodes in KG: user, item, category, among which 12 | users are placed ahead of other entities when indexing. 13 | CMD example: 14 | python main.py --model_name CFKG --emb_size 64 --margin 1 --include_attr 1 --lr 1e-4 --l2 1e-6 \ 15 | --dataset 'Grocery_and_Gourmet_Food' 16 | """ 17 | 18 | import torch 19 | import torch.nn as nn 20 | import numpy as np 21 | import pandas as pd 22 | 23 | from utils import utils 24 | from models.BaseModel import GeneralModel 25 | from helpers.KGReader import KGReader 26 | 27 | 28 | class CFKG(GeneralModel): 29 | reader = 'KGReader' 30 | runner = 'BaseRunner' 31 | extra_log_args = ['emb_size', 'margin', 'include_attr'] 32 | 33 | @staticmethod 34 | def parse_model_args(parser): 35 | parser.add_argument('--emb_size', type=int, default=64, 36 | help='Size of embedding vectors.') 37 | parser.add_argument('--margin', type=float, default=0, 38 | help='Margin in hinge loss.') 39 | return GeneralModel.parse_model_args(parser) 40 | 41 | def __init__(self, args, corpus: KGReader): 42 | super().__init__(args, corpus) 43 | self.emb_size = args.emb_size 44 | self.margin = args.margin 45 | self.relation_num = corpus.n_relations 46 | self.entity_num = corpus.n_entities 47 | self._define_params() 48 | self.apply(self.init_weights) 49 | 50 | def _define_params(self): 51 | self.e_embeddings = nn.Embedding(self.user_num + self.entity_num, self.emb_size) 52 | # ↑ user and entity embeddings, user first 53 | self.r_embeddings = nn.Embedding(self.relation_num, self.emb_size) 54 | # ↑ relation embedding: 0 is used for "buy" between users and items 55 | self.loss_function = nn.MarginRankingLoss(margin=self.margin) 56 | 57 | def forward(self, feed_dict): 58 | self.check_list = [] 59 | head_ids = feed_dict['head_id'] # [batch_size, -1] 60 | tail_ids = feed_dict['tail_id'] # [batch_size, -1] 61 | relation_ids = feed_dict['relation_id'] # [batch_size, -1] 62 | 63 | head_vectors = self.e_embeddings(head_ids) 64 | tail_vectors = self.e_embeddings(tail_ids) 65 | relation_vectors = self.r_embeddings(relation_ids) 66 | 67 | prediction = -((head_vectors + relation_vectors - tail_vectors)**2).sum(-1) 68 | return {'prediction': prediction.view(feed_dict['batch_size'], -1)} 69 | 70 | def loss(self, out_dict): 71 | predictions = out_dict['prediction'] 72 | batch_size = predictions.shape[0] 73 | pos_pred, neg_pred = predictions[:, :2].flatten(), predictions[:, 2:].flatten() 74 | target = torch.from_numpy(np.ones(batch_size * 2, dtype=np.float32)).to(self.device) 75 | loss = self.loss_function(pos_pred, neg_pred, target) 76 | return loss 77 | 78 | class Dataset(GeneralModel.Dataset): 79 | def __init__(self, model, corpus, phase): 80 | super().__init__(model, corpus, phase) 81 | if self.phase == 'train': 82 | interaction_df = pd.DataFrame({ 83 | 'head': self.data['user_id'], 84 | 'tail': self.data['item_id'], 85 | 'relation': np.zeros_like(self.data['user_id']) # "buy" relation 86 | }) 87 | self.data = utils.df_to_dict(pd.concat((self.corpus.relation_df, interaction_df), axis=0)) 88 | self.neg_heads = np.zeros(len(self), dtype=int) 89 | self.neg_tails = np.zeros(len(self), dtype=int) 90 | 91 | def _get_feed_dict(self, index): 92 | if self.phase == 'train': 93 | head, tail = self.data['head'][index], self.data['tail'][index] 94 | relation = self.data['relation'][index] 95 | head_id = np.array([head, head, head, self.neg_heads[index]]) 96 | tail_id = np.array([tail, tail, self.neg_tails[index], tail]) 97 | relation_id = np.array([relation] * 4) 98 | if relation > 0: # head is not a user 99 | head_id = head_id + self.corpus.n_users 100 | else: 101 | target_item = self.data['item_id'][index] 102 | if self.model.test_all: 103 | neg_items = np.arange(1, self.corpus.n_items) 104 | else: 105 | neg_items = self.data['neg_items'][index] 106 | tail_id = np.concatenate([[target_item], neg_items]) 107 | head_id = self.data['user_id'][index] * np.ones_like(tail_id) 108 | relation_id = np.zeros_like(tail_id) 109 | tail_id += self.corpus.n_users # tail must be a non-user entity 110 | 111 | feed_dict = {'head_id': head_id, 'tail_id': tail_id, 'relation_id': relation_id} 112 | return feed_dict 113 | 114 | def actions_before_epoch(self): 115 | for i in range(len(self)): 116 | head, tail, relation = self.data['head'][i], self.data['tail'][i], self.data['relation'][i] 117 | self.neg_tails[i] = np.random.randint(1, self.corpus.n_items) 118 | if relation == 0: # "buy" relation 119 | self.neg_heads[i] = np.random.randint(1, self.corpus.n_users) 120 | while self.neg_tails[i] in self.corpus.train_clicked_set[head]: 121 | self.neg_tails[i] = np.random.randint(1, self.corpus.n_items) 122 | while tail in self.corpus.train_clicked_set[self.neg_heads[i]]: 123 | self.neg_heads[i] = np.random.randint(1, self.corpus.n_users) 124 | else: 125 | self.neg_heads[i] = np.random.randint(1, self.corpus.n_entities) 126 | while (head, relation, self.neg_tails[i]) in self.corpus.triplet_set: 127 | self.neg_tails[i] = np.random.randint(1, self.corpus.n_entities) 128 | while (self.neg_heads[i], relation, tail) in self.corpus.triplet_set: 129 | self.neg_heads[i] = np.random.randint(1, self.corpus.n_entities) 130 | -------------------------------------------------------------------------------- /src/models/general/DirectAU.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Chenyang Wang 3 | # @Email : THUwangcy@gmail.com 4 | 5 | """ DirectAU 6 | Reference: 7 | "Towards Representation Alignment and Uniformity in Collaborative Filtering" 8 | Wang et al., KDD'2022. 9 | CMD example: 10 | python main.py --model_name DirectAU --dataset Grocery_and_Gourmet_Food \ 11 | --emb_size 64 --lr 1e-3 --l2 1e-6 --epoch 500 --gamma 0.3 12 | """ 13 | 14 | import torch 15 | import torch.nn as nn 16 | import torch.nn.functional as F 17 | 18 | from models.BaseModel import GeneralModel 19 | 20 | 21 | class DirectAU(GeneralModel): 22 | reader = 'BaseReader' 23 | runner = 'BaseRunner' 24 | extra_log_args = ['emb_size', 'gamma'] 25 | 26 | @staticmethod 27 | def parse_model_args(parser): 28 | parser.add_argument('--emb_size', type=int, default=64, 29 | help='Size of embedding vectors.') 30 | parser.add_argument('--gamma', type=float, default=1, 31 | help='Weight of the uniformity loss.') 32 | return GeneralModel.parse_model_args(parser) 33 | 34 | @staticmethod 35 | def init_weights(m): 36 | if 'Linear' in str(type(m)): 37 | nn.init.xavier_normal_(m.weight.data) 38 | if m.bias is not None: 39 | nn.init.normal_(m.bias.data) 40 | elif 'Embedding' in str(type(m)): 41 | nn.init.xavier_normal_(m.weight.data) 42 | 43 | def __init__(self, args, corpus): 44 | super().__init__(args, corpus) 45 | self.emb_size = args.emb_size 46 | self.gamma = args.gamma 47 | self._define_params() 48 | self.apply(self.init_weights) 49 | 50 | def _define_params(self): 51 | self.u_embeddings = nn.Embedding(self.user_num, self.emb_size) 52 | self.i_embeddings = nn.Embedding(self.item_num, self.emb_size) 53 | 54 | @staticmethod 55 | def alignment(x, y): 56 | x, y = F.normalize(x, dim=-1), F.normalize(y, dim=-1) 57 | return (x - y).norm(p=2, dim=1).pow(2).mean() 58 | 59 | @staticmethod 60 | def uniformity(x): 61 | x = F.normalize(x, dim=-1) 62 | return torch.pdist(x, p=2).pow(2).mul(-2).exp().mean().log() 63 | 64 | def forward(self, feed_dict): 65 | self.check_list = [] 66 | user, items = feed_dict['user_id'], feed_dict['item_id'] 67 | 68 | user_e = self.u_embeddings(user) 69 | item_e = self.i_embeddings(items) 70 | 71 | prediction = (user_e[:, None, :] * item_e).sum(dim=-1) # [batch_size, -1] 72 | out_dict = {'prediction': prediction} 73 | 74 | if feed_dict['phase'] == 'train': 75 | out_dict.update({ 76 | 'user_e': user_e, 77 | 'item_e': item_e.squeeze(1) 78 | }) 79 | 80 | return out_dict 81 | 82 | def loss(self, output): 83 | user_e, item_e = output['user_e'], output['item_e'] 84 | 85 | align = self.alignment(user_e, item_e) 86 | uniform = (self.uniformity(user_e) + self.uniformity(item_e)) / 2 87 | loss = align + self.gamma * uniform 88 | 89 | return loss 90 | 91 | class Dataset(GeneralModel.Dataset): 92 | # No need to sample negative items 93 | def actions_before_epoch(self): 94 | self.data['neg_items'] = [[] for _ in range(len(self))] 95 | -------------------------------------------------------------------------------- /src/models/general/LightGCN.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Chenyang Wang 3 | # @Email : THUwangcy@gmail.com 4 | 5 | import torch 6 | import numpy as np 7 | import torch.nn as nn 8 | import scipy.sparse as sp 9 | 10 | from models.BaseModel import GeneralModel 11 | from models.BaseImpressionModel import ImpressionModel 12 | 13 | class LightGCNBase(object): 14 | @staticmethod 15 | def parse_model_args(parser): 16 | parser.add_argument('--emb_size', type=int, default=64, 17 | help='Size of embedding vectors.') 18 | parser.add_argument('--n_layers', type=int, default=3, 19 | help='Number of LightGCN layers.') 20 | return parser 21 | 22 | @staticmethod 23 | def build_adjmat(user_count, item_count, train_mat, selfloop_flag=False): 24 | R = sp.dok_matrix((user_count, item_count), dtype=np.float32) 25 | for user in train_mat: 26 | for item in train_mat[user]: 27 | R[user, item] = 1 28 | R = R.tolil() 29 | 30 | adj_mat = sp.dok_matrix((user_count + item_count, user_count + item_count), dtype=np.float32) 31 | adj_mat = adj_mat.tolil() 32 | 33 | adj_mat[:user_count, user_count:] = R 34 | adj_mat[user_count:, :user_count] = R.T 35 | adj_mat = adj_mat.todok() 36 | 37 | def normalized_adj_single(adj): 38 | # D^-1/2 * A * D^-1/2 39 | rowsum = np.array(adj.sum(1)) + 1e-10 40 | 41 | d_inv_sqrt = np.power(rowsum, -0.5).flatten() 42 | d_inv_sqrt[np.isinf(d_inv_sqrt)] = 0. 43 | d_mat_inv_sqrt = sp.diags(d_inv_sqrt) 44 | 45 | bi_lap = d_mat_inv_sqrt.dot(adj).dot(d_mat_inv_sqrt) 46 | return bi_lap.tocoo() 47 | 48 | if selfloop_flag: 49 | norm_adj_mat = normalized_adj_single(adj_mat + sp.eye(adj_mat.shape[0])) 50 | else: 51 | norm_adj_mat = normalized_adj_single(adj_mat) 52 | 53 | return norm_adj_mat.tocsr() 54 | 55 | def _base_init(self, args, corpus): 56 | self.emb_size = args.emb_size 57 | self.n_layers = args.n_layers 58 | self.norm_adj = self.build_adjmat(corpus.n_users, corpus.n_items, corpus.train_clicked_set) 59 | self._base_define_params() 60 | self.apply(self.init_weights) 61 | 62 | def _base_define_params(self): 63 | self.encoder = LGCNEncoder(self.user_num, self.item_num, self.emb_size, self.norm_adj, self.n_layers) 64 | 65 | def forward(self, feed_dict): 66 | self.check_list = [] 67 | user, items = feed_dict['user_id'], feed_dict['item_id'] 68 | u_embed, i_embed = self.encoder(user, items) 69 | 70 | prediction = (u_embed[:, None, :] * i_embed).sum(dim=-1) # [batch_size, -1] 71 | u_v = u_embed.repeat(1,items.shape[1]).view(items.shape[0],items.shape[1],-1) 72 | i_v = i_embed 73 | return {'prediction': prediction.view(feed_dict['batch_size'], -1), 'u_v': u_v, 'i_v':i_v} 74 | 75 | class LightGCN(GeneralModel, LightGCNBase): 76 | reader = 'BaseReader' 77 | runner = 'BaseRunner' 78 | extra_log_args = ['emb_size', 'n_layers', 'batch_size'] 79 | 80 | @staticmethod 81 | def parse_model_args(parser): 82 | parser = LightGCNBase.parse_model_args(parser) 83 | return GeneralModel.parse_model_args(parser) 84 | 85 | def __init__(self, args, corpus): 86 | GeneralModel.__init__(self, args, corpus) 87 | self._base_init(args, corpus) 88 | 89 | def forward(self, feed_dict): 90 | out_dict = LightGCNBase.forward(self, feed_dict) 91 | return {'prediction': out_dict['prediction']} 92 | 93 | class LightGCNImpression(ImpressionModel, LightGCNBase): 94 | reader = 'ImpressionReader' 95 | runner = 'ImpressionRunner' 96 | extra_log_args = ['emb_size', 'n_layers', 'batch_size'] 97 | 98 | @staticmethod 99 | def parse_model_args(parser): 100 | parser = LightGCNBase.parse_model_args(parser) 101 | return ImpressionModel.parse_model_args(parser) 102 | 103 | def __init__(self, args, corpus): 104 | ImpressionModel.__init__(self, args, corpus) 105 | self._base_init(args, corpus) 106 | 107 | def forward(self, feed_dict): 108 | return LightGCNBase.forward(self, feed_dict) 109 | 110 | class LGCNEncoder(nn.Module): 111 | def __init__(self, user_count, item_count, emb_size, norm_adj, n_layers=3): 112 | super(LGCNEncoder, self).__init__() 113 | self.user_count = user_count 114 | self.item_count = item_count 115 | self.emb_size = emb_size 116 | self.layers = [emb_size] * n_layers 117 | self.norm_adj = norm_adj 118 | 119 | self.embedding_dict = self._init_model() 120 | self.sparse_norm_adj = self._convert_sp_mat_to_sp_tensor(self.norm_adj).cuda() 121 | 122 | def _init_model(self): 123 | initializer = nn.init.xavier_uniform_ 124 | embedding_dict = nn.ParameterDict({ 125 | 'user_emb': nn.Parameter(initializer(torch.empty(self.user_count, self.emb_size))), 126 | 'item_emb': nn.Parameter(initializer(torch.empty(self.item_count, self.emb_size))), 127 | }) 128 | return embedding_dict 129 | 130 | @staticmethod 131 | def _convert_sp_mat_to_sp_tensor(X): 132 | coo = X.tocoo() 133 | i = torch.LongTensor([coo.row, coo.col]) 134 | v = torch.from_numpy(coo.data).float() 135 | return torch.sparse.FloatTensor(i, v, coo.shape) 136 | 137 | def forward(self, users, items): 138 | ego_embeddings = torch.cat([self.embedding_dict['user_emb'], self.embedding_dict['item_emb']], 0) 139 | all_embeddings = [ego_embeddings] 140 | 141 | for k in range(len(self.layers)): 142 | ego_embeddings = torch.sparse.mm(self.sparse_norm_adj, ego_embeddings) 143 | all_embeddings += [ego_embeddings] 144 | 145 | all_embeddings = torch.stack(all_embeddings, dim=1) 146 | all_embeddings = torch.mean(all_embeddings, dim=1) 147 | 148 | user_all_embeddings = all_embeddings[:self.user_count, :] 149 | item_all_embeddings = all_embeddings[self.user_count:, :] 150 | 151 | user_embeddings = user_all_embeddings[users, :] 152 | item_embeddings = item_all_embeddings[items, :] 153 | 154 | return user_embeddings, item_embeddings 155 | -------------------------------------------------------------------------------- /src/models/general/NeuMF.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Chenyang Wang 3 | # @Email : THUwangcy@gmail.com 4 | 5 | """ NeuMF 6 | Reference: 7 | "Neural Collaborative Filtering" 8 | Xiangnan He et al., WWW'2017. 9 | Reference code: 10 | The authors' tensorflow implementation https://github.com/hexiangnan/neural_collaborative_filtering 11 | CMD example: 12 | python main.py --model_name NeuMF --emb_size 64 --layers '[64]' --lr 5e-4 --l2 1e-7 --dropout 0.2 \ 13 | --dataset 'Grocery_and_Gourmet_Food' 14 | """ 15 | 16 | import torch 17 | import torch.nn as nn 18 | 19 | from models.BaseModel import GeneralModel 20 | 21 | 22 | class NeuMF(GeneralModel): 23 | reader = 'BaseReader' 24 | runner = 'BaseRunner' 25 | extra_log_args = ['emb_size', 'layers'] 26 | 27 | @staticmethod 28 | def parse_model_args(parser): 29 | parser.add_argument('--emb_size', type=int, default=64, 30 | help='Size of embedding vectors.') 31 | parser.add_argument('--layers', type=str, default='[64]', 32 | help="Size of each layer.") 33 | return GeneralModel.parse_model_args(parser) 34 | 35 | def __init__(self, args, corpus): 36 | super().__init__(args, corpus) 37 | self.emb_size = args.emb_size 38 | self.layers = eval(args.layers) 39 | self._define_params() 40 | self.apply(self.init_weights) 41 | 42 | def _define_params(self): 43 | self.mf_u_embeddings = nn.Embedding(self.user_num, self.emb_size) 44 | self.mf_i_embeddings = nn.Embedding(self.item_num, self.emb_size) 45 | self.mlp_u_embeddings = nn.Embedding(self.user_num, self.emb_size) 46 | self.mlp_i_embeddings = nn.Embedding(self.item_num, self.emb_size) 47 | 48 | self.mlp = nn.ModuleList([]) 49 | pre_size = 2 * self.emb_size 50 | for i, layer_size in enumerate(self.layers): 51 | self.mlp.append(nn.Linear(pre_size, layer_size)) 52 | pre_size = layer_size 53 | self.dropout_layer = nn.Dropout(p=self.dropout) 54 | self.prediction = nn.Linear(pre_size + self.emb_size, 1, bias=False) 55 | 56 | def forward(self, feed_dict): 57 | self.check_list = [] 58 | u_ids = feed_dict['user_id'] # [batch_size] 59 | i_ids = feed_dict['item_id'] # [batch_size, -1] 60 | 61 | u_ids = u_ids.unsqueeze(-1).repeat((1, i_ids.shape[1])) # [batch_size, -1] 62 | 63 | mf_u_vectors = self.mf_u_embeddings(u_ids) 64 | mf_i_vectors = self.mf_i_embeddings(i_ids) 65 | mlp_u_vectors = self.mlp_u_embeddings(u_ids) 66 | mlp_i_vectors = self.mlp_i_embeddings(i_ids) 67 | 68 | mf_vector = mf_u_vectors * mf_i_vectors 69 | mlp_vector = torch.cat([mlp_u_vectors, mlp_i_vectors], dim=-1) 70 | for layer in self.mlp: 71 | mlp_vector = layer(mlp_vector).relu() 72 | mlp_vector = self.dropout_layer(mlp_vector) 73 | 74 | output_vector = torch.cat([mf_vector, mlp_vector], dim=-1) 75 | prediction = self.prediction(output_vector) 76 | return {'prediction': prediction.view(feed_dict['batch_size'], -1)} 77 | -------------------------------------------------------------------------------- /src/models/general/POP.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import torch 4 | import numpy as np 5 | 6 | from models.BaseModel import GeneralModel 7 | 8 | 9 | class POP(GeneralModel): 10 | """ 11 | Recommendation according to item's popularity. 12 | Should run with --train 0 13 | """ 14 | def __init__(self, args, corpus): 15 | super().__init__(args, corpus) 16 | self.popularity = np.zeros(corpus.n_items) 17 | for i in corpus.data_df['train']['item_id'].values: 18 | self.popularity[i] += 1 19 | 20 | def forward(self, feed_dict): 21 | self.check_list = [] 22 | i_ids = feed_dict['item_id'] # [batch_size, -1] 23 | prediction = self.popularity[i_ids.cpu().data.numpy()] 24 | prediction = torch.from_numpy(prediction).to(self.device) 25 | return {'prediction': prediction.view(feed_dict['batch_size'], -1)} 26 | -------------------------------------------------------------------------------- /src/models/general/__init__.py: -------------------------------------------------------------------------------- 1 | from os.path import dirname, basename, isfile, join 2 | import glob 3 | 4 | modules = glob.glob(join(dirname(__file__), "*.py")) 5 | __all__ = [ 6 | basename(f)[:-3] for f in modules if isfile(f) and not f.endswith('__init__.py') 7 | ] 8 | -------------------------------------------------------------------------------- /src/models/reranker/PRM.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Hanyu Li 3 | # @Email : hanyu-li23@mails.tsinghua.edu.cn 4 | 5 | """ PRM 6 | Reference: 7 | "Personalized Re-ranking for Recommendation" 8 | Pei et al., RecSys'2019. 9 | """ 10 | 11 | import torch 12 | import torch.nn as nn 13 | from models.BaseRerankerModel import RerankModel 14 | from models.BaseRerankerModel import RerankSeqModel 15 | from models.general import * 16 | from models.sequential import * 17 | from models.developing import * 18 | 19 | class LearnedPositionEncoding(nn.Embedding): 20 | def __init__(self,d_model, dropout = 0,max_len = 50): 21 | super().__init__(max_len, d_model) 22 | self.dropout = nn.Dropout(p = dropout) 23 | 24 | def forward(self, x): 25 | weight = self.weight.data.unsqueeze(1) 26 | x = x + weight[:x.size(0),:] 27 | return self.dropout(x) 28 | 29 | class PRMBase(object): 30 | @staticmethod 31 | def parse_model_args(parser): 32 | parser.add_argument('--emb_size', type=int, default=64, 33 | help='Size of item embedding vectors.') 34 | parser.add_argument('--n_blocks', type=int, default=4, 35 | help='num of blocks of MSAB/IMSAB') 36 | parser.add_argument('--num_heads', type=int, default=4, 37 | help='Number of attention heads.') 38 | parser.add_argument('--num_hidden_unit', type=int, default=64, 39 | help='Number of hidden units in Transformer layer.') 40 | return parser 41 | 42 | def _base_init(self, args, corpus): 43 | self.args = args 44 | self.emb_size = args.emb_size 45 | self.n_blocks = args.n_blocks 46 | self.num_heads = args.num_heads 47 | self.num_hidden_unit = args.num_hidden_unit 48 | self.positionafter = 0 49 | 50 | self.dropout = args.dropout 51 | 52 | self.corpus=corpus 53 | 54 | self._base_define_params() 55 | self.apply(self.init_weights) 56 | 57 | def _base_define_params(self): 58 | self.i_embeddings = nn.Embedding(self.item_num, self.emb_size) 59 | if self.positionafter == 0: 60 | self.ordinal_position_embedding = nn.Embedding(self.train_max_neg_item + self.train_max_pos_item, self.emb_size + self.ranker_emb_size * 2) # reranker and ranker embedding 61 | else: 62 | self.ordinal_position_embedding=nn.Embedding(self.train_max_neg_item+self.train_max_pos_item,self.num_hidden_unit) 63 | self.rFF0 = nn.Linear(self.emb_size + self.ranker_emb_size * 2, self.num_hidden_unit, bias=True) 64 | self.encoder = nn.ModuleList([nn.TransformerEncoderLayer(d_model=self.num_hidden_unit, nhead=self.num_heads, dim_feedforward=128, dropout=self.dropout) for _ in range(self.n_blocks)]) 65 | self.rFF1 = nn.Linear(self.num_hidden_unit, 1, bias=True) 66 | 67 | def forward(self, feed_dict): 68 | batch_size = feed_dict['item_id'].shape[0] 69 | #history_max = feed_dict['history_items'].shape[1] 70 | 71 | i_ids = feed_dict['item_id'] # [batch_size, sample_num] 72 | #u_ids = feed_dict['user_id'] # [batch_size, sample_num] Here this should be switched to PV vector, which is from the pretrained model and represents the relationship between u and i 73 | 74 | i_vectors = self.i_embeddings(i_ids) # [batch_size, sample_num, emb_size] 75 | u_vectors = torch.cat([feed_dict['u_v'],feed_dict['i_v']],dim=2)#pv, consist of sequence vector and candidate embedding of the base ranker model, attaching the sequence embedding to every candidate emb [batch_size, sequence_emb+item_emb] 76 | #score = feed_dict['scores'] 77 | 78 | di = torch.cat((i_vectors,u_vectors),dim=2) 79 | pi = self.ordinal_position_embedding(feed_dict['position']) #learnable position encoding 80 | 81 | if self.positionafter==0: 82 | xi = di+pi 83 | xi = self.rFF0(xi) #reshape dimension to num_hidden_unit 84 | else: 85 | xi = self.rFF0(di) #reshape dimension to num_hidden_unit 86 | xi = xi+pi 87 | 88 | padding_mask = feed_dict['padding_mask'] 89 | 90 | xi = torch.transpose(xi, 0, 1) 91 | for block in self.encoder: 92 | xi = block(xi, None, padding_mask) 93 | 94 | prediction = self.rFF1(xi) 95 | prediction = torch.transpose(prediction, 0, 1) 96 | 97 | return {'prediction': prediction.view(batch_size, -1)} 98 | 99 | class PRMGeneral(RerankModel, PRMBase): 100 | reader = 'ImpressionReader' 101 | runner = 'ImpressionRunner' 102 | 103 | @staticmethod 104 | def parse_model_args(parser): 105 | parser = PRMBase.parse_model_args(parser) 106 | return RerankModel.parse_model_args(parser) 107 | 108 | def __init__(self, args, corpus): 109 | RerankModel.__init__(self, args, corpus) 110 | self._base_init(args, corpus) 111 | 112 | def forward(self, feed_dict): 113 | return PRMBase.forward(self, feed_dict) 114 | 115 | class PRMSequential(RerankSeqModel, PRMBase): 116 | reader = 'ImpressionSeqReader' 117 | runner = 'ImpressionRunner' 118 | 119 | @staticmethod 120 | def parse_model_args(parser): 121 | parser = PRMBase.parse_model_args(parser) 122 | return RerankSeqModel.parse_model_args(parser) 123 | 124 | def __init__(self, args, corpus): 125 | RerankSeqModel.__init__(self, args, corpus) 126 | self._base_init(args, corpus) 127 | 128 | def forward(self, feed_dict): 129 | return PRMBase.forward(self, feed_dict) -------------------------------------------------------------------------------- /src/models/reranker/__init__.py: -------------------------------------------------------------------------------- 1 | from os.path import dirname, basename, isfile, join 2 | import glob 3 | 4 | modules = glob.glob(join(dirname(__file__), "*.py")) 5 | __all__ = [ 6 | basename(f)[:-3] for f in modules if isfile(f) and not f.endswith('__init__.py') 7 | ] 8 | -------------------------------------------------------------------------------- /src/models/sequential/Caser.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Chenyang Wang 3 | # @Email : THUwangcy@gmail.com 4 | 5 | """ Caser 6 | Reference: 7 | "Personalized Top-N Sequential Recommendation via Convolutional Sequence Embedding" 8 | Jiaxi Tang et al., WSDM'2018. 9 | Reference code: 10 | https://github.com/graytowne/caser_pytorch 11 | Note: 12 | We use a maximum of L (instead of history_max) horizontal filters to prevent excessive CNN layers. 13 | Besides, to keep consistent with other sequential models, we do not use the sliding window to generate 14 | training instances in the paper, and set the parameter T as 1. 15 | CMD example: 16 | python main.py --model_name Caser --emb_size 64 --L 5 --num_horizon 64 --num_vertical 32 --lr 1e-3 --l2 1e-4 \ 17 | --history_max 20 --dataset 'Grocery_and_Gourmet_Food' 18 | """ 19 | 20 | import torch 21 | from torch import nn 22 | import torch.nn.functional as F 23 | 24 | from models.BaseModel import SequentialModel 25 | 26 | 27 | class Caser(SequentialModel): 28 | reader = 'SeqReader' 29 | runner = 'BaseRunner' 30 | extra_log_args = ['emb_size', 'num_horizon', 'num_vertical', 'L'] 31 | 32 | @staticmethod 33 | def parse_model_args(parser): 34 | parser.add_argument('--emb_size', type=int, default=64, 35 | help='Size of embedding vectors.') 36 | parser.add_argument('--num_horizon', type=int, default=16, 37 | help='Number of horizon convolution kernels.') 38 | parser.add_argument('--num_vertical', type=int, default=8, 39 | help='Number of vertical convolution kernels.') 40 | parser.add_argument('--L', type=int, default=4, 41 | help='Union window size.') 42 | return SequentialModel.parse_model_args(parser) 43 | 44 | def __init__(self, args, corpus): 45 | super().__init__(args, corpus) 46 | self.emb_size = args.emb_size 47 | self.max_his = args.history_max 48 | self.num_horizon = args.num_horizon 49 | self.num_vertical = args.num_vertical 50 | self.l = args.L 51 | assert self.l <= self.max_his # use L instead of max_his to avoid excessive conv_h 52 | self._define_params() 53 | self.apply(self.init_weights) 54 | 55 | def _define_params(self): 56 | self.u_embeddings = nn.Embedding(self.user_num, self.emb_size) 57 | self.i_embeddings = nn.Embedding(self.item_num, self.emb_size, padding_idx=0) 58 | lengths = [i + 1 for i in range(self.l)] 59 | self.conv_h = nn.ModuleList( 60 | [nn.Conv2d(in_channels=1, out_channels=self.num_horizon, kernel_size=(i, self.emb_size)) for i in lengths]) 61 | self.conv_v = nn.Conv2d(in_channels=1, out_channels=self.num_vertical, kernel_size=(self.max_his, 1)) 62 | 63 | self.fc_dim_h = self.num_horizon * len(lengths) 64 | self.fc_dim_v = self.num_vertical * self.emb_size 65 | fc_dim_in = self.fc_dim_v + self.fc_dim_h 66 | self.fc = nn.Linear(fc_dim_in, self.emb_size) 67 | self.out = nn.Linear(self.emb_size * 2, self.emb_size) 68 | 69 | 70 | def forward(self, feed_dict): 71 | self.check_list = [] 72 | u_ids = feed_dict['user_id'] 73 | i_ids = feed_dict['item_id'] # [batch_size, -1] 74 | history = feed_dict['history_items'] # [batch_size, history_max] 75 | batch_size, seq_len = history.shape 76 | 77 | pad_len = self.max_his - seq_len 78 | history = F.pad(history, [0, pad_len]) 79 | his_vectors = self.i_embeddings(history).unsqueeze(1) # [batch_size, 1, history_max, emb_size] 80 | 81 | # Convolution Layers 82 | out, out_h, out_v = None, None, None 83 | # vertical conv layer 84 | if self.num_vertical > 0: 85 | out_v = self.conv_v(his_vectors) 86 | out_v = out_v.view(-1, self.fc_dim_v) # prepare for fully connect 87 | # horizontal conv layer 88 | out_hs = list() 89 | if self.num_horizon > 0: 90 | for conv in self.conv_h: 91 | conv_out = conv(his_vectors).squeeze(3).relu() 92 | pool_out = F.max_pool1d(conv_out, conv_out.size(2)).squeeze(2) 93 | out_hs.append(pool_out) 94 | out_h = torch.cat(out_hs, 1) # prepare for fully connect 95 | 96 | # Fully-connected Layers 97 | user_vector = self.u_embeddings(u_ids) 98 | z = self.fc(torch.cat([out_v, out_h], 1)).relu() 99 | his_vector = self.out(torch.cat([z, user_vector], 1)) 100 | 101 | i_vectors = self.i_embeddings(i_ids) 102 | prediction = (his_vector[:, None, :] * i_vectors).sum(-1) 103 | return {'prediction': prediction.view(batch_size, -1)} -------------------------------------------------------------------------------- /src/models/sequential/ComiRec.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Chenyang Wang 3 | # @Email : THUwangcy@gmail.com 4 | 5 | """ ComiRec 6 | Reference: 7 | "Controllable Multi-Interest Framework for Recommendation" 8 | Cen et al., KDD'2020. 9 | CMD example: 10 | python main.py --model_name ComiRec --emb_size 64 --lr 1e-3 --l2 1e-6 --attn_size 8 --K 4 --add_pos 1 \ 11 | --history_max 20 --dataset 'Grocery_and_Gourmet_Food' 12 | """ 13 | 14 | import torch 15 | import torch.nn as nn 16 | import numpy as np 17 | 18 | from models.BaseModel import SequentialModel 19 | from utils import layers 20 | 21 | 22 | class ComiRec(SequentialModel): 23 | reader = 'SeqReader' 24 | runner = 'BaseRunner' 25 | extra_log_args = ['emb_size', 'attn_size', 'K'] 26 | 27 | @staticmethod 28 | def parse_model_args(parser): 29 | parser.add_argument('--emb_size', type=int, default=64, 30 | help='Size of embedding vectors.') 31 | parser.add_argument('--attn_size', type=int, default=8, 32 | help='Size of attention vectors.') 33 | parser.add_argument('--K', type=int, default=2, 34 | help='Number of hidden intent.') 35 | parser.add_argument('--add_pos', type=int, default=1, 36 | help='Whether add position embedding.') 37 | return SequentialModel.parse_model_args(parser) 38 | 39 | def __init__(self, args, corpus): 40 | super().__init__(args, corpus) 41 | self.emb_size = args.emb_size 42 | self.attn_size = args.attn_size 43 | self.K = args.K 44 | self.add_pos = args.add_pos 45 | self.max_his = args.history_max 46 | self.len_range = torch.from_numpy(np.arange(self.max_his)).to(self.device) 47 | self._define_params() 48 | self.apply(self.init_weights) 49 | 50 | def _define_params(self): 51 | self.i_embeddings = nn.Embedding(self.item_num, self.emb_size) 52 | if self.add_pos: 53 | self.p_embeddings = nn.Embedding(self.max_his + 1, self.emb_size) 54 | self.W1 = nn.Linear(self.emb_size, self.attn_size) 55 | self.W2 = nn.Linear(self.attn_size, self.K) 56 | 57 | def forward(self, feed_dict): 58 | self.check_list = [] 59 | i_ids = feed_dict['item_id'] # [batch_size, -1] 60 | history = feed_dict['history_items'] # [batch_size, history_max] 61 | lengths = feed_dict['lengths'] # [batch_size] 62 | batch_size, seq_len = history.shape 63 | 64 | valid_his = (history > 0).long() 65 | his_vectors = self.i_embeddings(history) 66 | 67 | if self.add_pos: 68 | position = (lengths[:, None] - self.len_range[None, :seq_len]) * valid_his 69 | pos_vectors = self.p_embeddings(position) 70 | his_pos_vectors = his_vectors + pos_vectors 71 | else: 72 | his_pos_vectors = his_vectors 73 | 74 | # Self-attention 75 | attn_score = self.W2(self.W1(his_pos_vectors).tanh()) # bsz, his_max, K 76 | attn_score = attn_score.masked_fill(valid_his.unsqueeze(-1) == 0, -np.inf) 77 | attn_score = attn_score.transpose(-1, -2) # bsz, K, his_max 78 | attn_score = (attn_score - attn_score.max()).softmax(dim=-1) 79 | attn_score = attn_score.masked_fill(torch.isnan(attn_score), 0) 80 | interest_vectors = (his_vectors[:, None, :, :] * attn_score[:, :, :, None]).sum(-2) # bsz, K, emb 81 | 82 | i_vectors = self.i_embeddings(i_ids) 83 | if feed_dict['phase'] == 'train': 84 | target_vector = i_vectors[:, 0] # bsz, emb 85 | target_pred = (interest_vectors * target_vector[:, None, :]).sum(-1) # bsz, K 86 | idx_select = target_pred.max(-1)[1] # bsz 87 | user_vector = interest_vectors[torch.arange(batch_size), idx_select, :] # bsz, emb 88 | prediction = (user_vector[:, None, :] * i_vectors).sum(-1) 89 | else: 90 | prediction = (interest_vectors[:, None, :, :] * i_vectors[:, :, None, :]).sum(-1) # bsz, -1, K 91 | prediction = prediction.max(-1)[0] # bsz, -1 92 | 93 | return {'prediction': prediction.view(batch_size, -1)} 94 | -------------------------------------------------------------------------------- /src/models/sequential/FPMC.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Chenyang Wang 3 | # @Email : THUwangcy@gmail.com 4 | 5 | """ FPMC 6 | Reference: 7 | "Factorizing Personalized Markov Chains for Next-Basket Recommendation" 8 | Rendle et al., WWW'2010. 9 | CMD example: 10 | python main.py --model_name FPMC --emb_size 64 --lr 1e-3 --l2 1e-6 --history_max 20 \ 11 | --dataset 'Grocery_and_Gourmet_Food' 12 | """ 13 | 14 | import torch 15 | import torch.nn as nn 16 | import numpy as np 17 | 18 | from models.BaseModel import SequentialModel 19 | 20 | 21 | class FPMC(SequentialModel): 22 | reader = 'SeqReader' 23 | runner = 'BaseRunner' 24 | extra_log_args = ['emb_size'] 25 | 26 | @staticmethod 27 | def parse_model_args(parser): 28 | parser.add_argument('--emb_size', type=int, default=64, 29 | help='Size of embedding vectors.') 30 | return SequentialModel.parse_model_args(parser) 31 | 32 | def __init__(self, args, corpus): 33 | super().__init__(args, corpus) 34 | self.emb_size = args.emb_size 35 | self._define_params() 36 | self.apply(self.init_weights) 37 | 38 | def _define_params(self): 39 | self.ui_embeddings = nn.Embedding(self.user_num, self.emb_size) 40 | self.iu_embeddings = nn.Embedding(self.item_num, self.emb_size) 41 | self.li_embeddings = nn.Embedding(self.item_num, self.emb_size) 42 | self.il_embeddings = nn.Embedding(self.item_num, self.emb_size) 43 | 44 | def forward(self, feed_dict): 45 | self.check_list = [] 46 | u_id = feed_dict['user_id'] # [batch_size] 47 | i_ids = feed_dict['item_id'] # [batch_size, -1] 48 | li_id = feed_dict['last_item_id'] # [batch_size] 49 | 50 | ui_vectors = self.ui_embeddings(u_id) 51 | iu_vectors = self.iu_embeddings(i_ids) 52 | li_vectors = self.li_embeddings(li_id) 53 | il_vectors = self.il_embeddings(i_ids) 54 | 55 | prediction = (ui_vectors[:, None, :] * iu_vectors).sum(-1) + (li_vectors[:, None, :] * il_vectors).sum(-1) 56 | return {'prediction': prediction.view(feed_dict['batch_size'], -1)} 57 | 58 | class Dataset(SequentialModel.Dataset): 59 | def _get_feed_dict(self, index): 60 | user_id, target_item = self.data['user_id'][index], self.data['item_id'][index] 61 | if self.phase != 'train' and self.model.test_all: 62 | neg_items = np.arange(1, self.corpus.n_items) 63 | else: 64 | neg_items = self.data['neg_items'][index] 65 | item_ids = np.concatenate([[target_item], neg_items]).astype(int) 66 | pos = self.data['position'][index] 67 | last_item_id = self.corpus.user_his[user_id][pos - 1][0] 68 | feed_dict = { 69 | 'user_id': user_id, 70 | 'item_id': item_ids, 71 | 'last_item_id': last_item_id 72 | } 73 | return feed_dict 74 | -------------------------------------------------------------------------------- /src/models/sequential/GRU4Rec.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Chenyang Wang 3 | # @Email : THUwangcy@gmail.com 4 | 5 | """ GRU4Rec 6 | Reference: 7 | "Session-based Recommendations with Recurrent Neural Networks" 8 | Hidasi et al., ICLR'2016. 9 | CMD example: 10 | python main.py --model_name GRU4Rec --emb_size 64 --hidden_size 128 --lr 1e-3 --l2 1e-4 --history_max 20 \ 11 | --dataset 'Grocery_and_Gourmet_Food' 12 | """ 13 | 14 | import torch 15 | import torch.nn as nn 16 | 17 | from models.BaseModel import SequentialModel 18 | from models.BaseImpressionModel import ImpressionSeqModel 19 | 20 | class GRU4RecBase(object): 21 | @staticmethod 22 | def parse_model_args(parser): 23 | parser.add_argument('--emb_size', type=int, default=64, 24 | help='Size of embedding vectors.') 25 | parser.add_argument('--hidden_size', type=int, default=64, 26 | help='Size of hidden vectors in GRU.') 27 | return parser 28 | 29 | def _base_init(self, args, corpus): 30 | self.emb_size = args.emb_size 31 | self.hidden_size = args.hidden_size 32 | self._base_define_params() 33 | self.apply(self.init_weights) 34 | 35 | def _base_define_params(self): 36 | self.i_embeddings = nn.Embedding(self.item_num, self.emb_size) 37 | self.rnn = nn.GRU(input_size=self.emb_size, hidden_size=self.hidden_size, batch_first=True) 38 | # self.pred_embeddings = nn.Embedding(self.item_num, self.hidden_size) 39 | self.out = nn.Linear(self.hidden_size, self.emb_size) 40 | 41 | def forward(self, feed_dict): 42 | self.check_list = [] 43 | i_ids = feed_dict['item_id'] # [batch_size, -1] 44 | history = feed_dict['history_items'] # [batch_size, history_max] 45 | lengths = feed_dict['lengths'] # [batch_size] 46 | 47 | his_vectors = self.i_embeddings(history) 48 | 49 | # Sort and Pack 50 | sort_his_lengths, sort_idx = torch.topk(lengths, k=len(lengths)) 51 | sort_his_vectors = his_vectors.index_select(dim=0, index=sort_idx) 52 | history_packed = torch.nn.utils.rnn.pack_padded_sequence( 53 | sort_his_vectors, sort_his_lengths.cpu(), batch_first=True) 54 | 55 | # RNN 56 | output, hidden = self.rnn(history_packed, None) 57 | 58 | # Unsort 59 | unsort_idx = torch.topk(sort_idx, k=len(lengths), largest=False)[1] 60 | rnn_vector = hidden[-1].index_select(dim=0, index=unsort_idx) 61 | 62 | # Predicts 63 | # pred_vectors = self.pred_embeddings(i_ids) 64 | pred_vectors = self.i_embeddings(i_ids) 65 | rnn_vector = self.out(rnn_vector) 66 | prediction = (rnn_vector[:, None, :] * pred_vectors).sum(-1) 67 | 68 | u_v = rnn_vector.repeat(1,i_ids.shape[1]).view(i_ids.shape[0],i_ids.shape[1],-1) 69 | i_v = pred_vectors 70 | 71 | return {'prediction': prediction.view(feed_dict['batch_size'], -1), 72 | 'u_v': u_v, 'i_v': i_v} 73 | 74 | class GRU4Rec(SequentialModel, GRU4RecBase): 75 | reader = 'SeqReader' 76 | runner = 'BaseRunner' 77 | extra_log_args = ['emb_size', 'hidden_size'] 78 | 79 | @staticmethod 80 | def parse_model_args(parser): 81 | parser = GRU4RecBase.parse_model_args(parser) 82 | return SequentialModel.parse_model_args(parser) 83 | 84 | def __init__(self, args, corpus): 85 | SequentialModel.__init__(self, args, corpus) 86 | self._base_init(args, corpus) 87 | 88 | def forward(self, feed_dict): 89 | out_dict = GRU4RecBase.forward(self, feed_dict) 90 | return {'prediction': out_dict['prediction']} 91 | 92 | class GRU4RecImpression(ImpressionSeqModel, GRU4RecBase): 93 | reader = 'ImpressionSeqReader' 94 | runner = 'ImpressionRunner' 95 | extra_log_args = ['emb_size', 'hidden_size'] 96 | 97 | @staticmethod 98 | def parse_model_args(parser): 99 | parser = GRU4RecBase.parse_model_args(parser) 100 | return ImpressionSeqModel.parse_model_args(parser) 101 | 102 | def __init__(self, args, corpus): 103 | ImpressionSeqModel.__init__(self, args, corpus) 104 | self._base_init(args, corpus) 105 | 106 | def forward(self, feed_dict): 107 | return GRU4RecBase.forward(self, feed_dict) -------------------------------------------------------------------------------- /src/models/sequential/NARM.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Chenyang Wang 3 | # @Email : THUwangcy@gmail.com 4 | 5 | """ NARM 6 | Reference: 7 | "Neural Attentive Session-based Recommendation" 8 | Jing Li et al., CIKM'2017. 9 | CMD example: 10 | python main.py --model_name NARM --emb_size 64 --hidden_size 100 --attention_size 4 --lr 1e-3 --l2 1e-4 \ 11 | --history_max 20 --dataset 'Grocery_and_Gourmet_Food' 12 | """ 13 | 14 | import torch 15 | import torch.nn as nn 16 | 17 | from models.BaseModel import SequentialModel 18 | 19 | 20 | class NARM(SequentialModel): 21 | reader = 'SeqReader' 22 | runner = 'BaseRunner' 23 | extra_log_args = ['emb_size', 'hidden_size', 'attention_size'] 24 | 25 | @staticmethod 26 | def parse_model_args(parser): 27 | parser.add_argument('--emb_size', type=int, default=64, 28 | help='Size of embedding vectors.') 29 | parser.add_argument('--hidden_size', type=int, default=100, 30 | help='Size of hidden vectors in GRU.') 31 | parser.add_argument('--attention_size', type=int, default=50, 32 | help='Size of attention hidden space.') 33 | return SequentialModel.parse_model_args(parser) 34 | 35 | def __init__(self, args, corpus): 36 | super().__init__(args, corpus) 37 | self.emb_size = args.emb_size 38 | self.hidden_size = args.hidden_size 39 | self.attention_size = args.attention_size 40 | self._define_params() 41 | self.apply(self.init_weights) 42 | 43 | def _define_params(self): 44 | self.i_embeddings = nn.Embedding(self.item_num, self.emb_size) 45 | self.encoder_g = nn.GRU(input_size=self.emb_size, hidden_size=self.hidden_size, batch_first=True) 46 | self.encoder_l = nn.GRU(input_size=self.emb_size, hidden_size=self.hidden_size, batch_first=True) 47 | self.A1 = nn.Linear(self.hidden_size, self.attention_size, bias=False) 48 | self.A2 = nn.Linear(self.hidden_size, self.attention_size, bias=False) 49 | self.attention_out = nn.Linear(self.attention_size, 1, bias=False) 50 | self.out = nn.Linear(2 * self.hidden_size, self.emb_size, bias=False) 51 | 52 | def forward(self, feed_dict): 53 | self.check_list = [] 54 | i_ids = feed_dict['item_id'] # [batch_size, -1] 55 | history = feed_dict['history_items'] # [batch_size, history_max] 56 | lengths = feed_dict['lengths'] # [batch_size] 57 | 58 | # Embedding Layer 59 | i_vectors = self.i_embeddings(i_ids) 60 | his_vectors = self.i_embeddings(history) 61 | 62 | # Encoding Layer 63 | sort_his_lengths, sort_idx = torch.topk(lengths, k=len(lengths)) 64 | sort_his_vectors = his_vectors.index_select(dim=0, index=sort_idx) 65 | history_packed = nn.utils.rnn.pack_padded_sequence(sort_his_vectors, sort_his_lengths.cpu(), batch_first=True) 66 | _, hidden_g = self.encoder_g(history_packed, None) 67 | output_l, hidden_l = self.encoder_l(history_packed, None) 68 | output_l, _ = torch.nn.utils.rnn.pad_packed_sequence(output_l, batch_first=True) 69 | unsort_idx = torch.topk(sort_idx, k=len(lengths), largest=False)[1] 70 | output_l = output_l.index_select(dim=0, index=unsort_idx) # [batch_size, history_max, emb_size] 71 | hidden_g = hidden_g[-1].index_select(dim=0, index=unsort_idx) # [batch_size, emb_size] 72 | 73 | # Attention Layer 74 | attention_g = self.A1(hidden_g) 75 | attention_l = self.A2(output_l) 76 | attention_value = self.attention_out((attention_g[:, None, :] + attention_l).sigmoid()) 77 | mask = (history > 0).unsqueeze(-1) 78 | attention_value = attention_value.masked_fill(mask == 0, 0) 79 | c_l = (attention_value * output_l).sum(1) 80 | 81 | # Prediction Layer 82 | pred_vector = self.out(torch.cat((hidden_g, c_l), dim=1)) 83 | prediction = (pred_vector[:, None, :] * i_vectors).sum(dim=-1) 84 | return {'prediction': prediction.view(feed_dict['batch_size'], -1)} 85 | -------------------------------------------------------------------------------- /src/models/sequential/SASRec.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Chenyang Wang 3 | # @Email : THUwangcy@gmail.com 4 | 5 | """ SASRec 6 | Reference: 7 | "Self-attentive Sequential Recommendation" 8 | Kang et al., IEEE'2018. 9 | Note: 10 | When incorporating position embedding, we make the position index start from the most recent interaction. 11 | """ 12 | 13 | import torch 14 | import torch.nn as nn 15 | import numpy as np 16 | 17 | from models.BaseModel import SequentialModel 18 | from models.BaseImpressionModel import ImpressionSeqModel 19 | from utils import layers 20 | 21 | class SASRecBase(object): 22 | @staticmethod 23 | def parse_model_args(parser): 24 | parser.add_argument('--emb_size', type=int, default=64, 25 | help='Size of embedding vectors.') 26 | parser.add_argument('--num_layers', type=int, default=1, 27 | help='Number of self-attention layers.') 28 | parser.add_argument('--num_heads', type=int, default=4, 29 | help='Number of attention heads.') 30 | return parser 31 | 32 | def _base_init(self, args, corpus): 33 | self.emb_size = args.emb_size 34 | self.max_his = args.history_max 35 | self.num_layers = args.num_layers 36 | self.num_heads = args.num_heads 37 | self.len_range = torch.from_numpy(np.arange(self.max_his)).to(self.device) 38 | self._base_define_params() 39 | self.apply(self.init_weights) 40 | 41 | def _base_define_params(self): 42 | self.i_embeddings = nn.Embedding(self.item_num, self.emb_size) 43 | self.p_embeddings = nn.Embedding(self.max_his + 1, self.emb_size) 44 | 45 | self.transformer_block = nn.ModuleList([ 46 | layers.TransformerLayer(d_model=self.emb_size, d_ff=self.emb_size, n_heads=self.num_heads, 47 | dropout=self.dropout, kq_same=False) 48 | for _ in range(self.num_layers) 49 | ]) 50 | 51 | def forward(self, feed_dict): 52 | self.check_list = [] 53 | i_ids = feed_dict['item_id'] # [batch_size, -1] 54 | history = feed_dict['history_items'] # [batch_size, history_max] 55 | lengths = feed_dict['lengths'] # [batch_size] 56 | batch_size, seq_len = history.shape 57 | 58 | valid_his = (history > 0).long() 59 | his_vectors = self.i_embeddings(history) 60 | 61 | # Position embedding 62 | # lengths: [4, 2, 5] 63 | # position: [[4, 3, 2, 1, 0], [2, 1, 0, 0, 0], [5, 4, 3, 2, 1]] 64 | position = (lengths[:, None] - self.len_range[None, :seq_len]) * valid_his 65 | pos_vectors = self.p_embeddings(position) 66 | his_vectors = his_vectors + pos_vectors 67 | 68 | # Self-attention 69 | causality_mask = np.tril(np.ones((1, 1, seq_len, seq_len), dtype=np.int)) 70 | attn_mask = torch.from_numpy(causality_mask).to(self.device) 71 | # attn_mask = valid_his.view(batch_size, 1, 1, seq_len) 72 | for block in self.transformer_block: 73 | his_vectors = block(his_vectors, attn_mask) 74 | his_vectors = his_vectors * valid_his[:, :, None].float() 75 | 76 | his_vector = his_vectors[torch.arange(batch_size), lengths - 1, :] 77 | # his_vector = his_vectors.sum(1) / lengths[:, None].float() 78 | # ↑ average pooling is shown to be more effective than the most recent embedding 79 | 80 | i_vectors = self.i_embeddings(i_ids) 81 | prediction = (his_vector[:, None, :] * i_vectors).sum(-1) 82 | 83 | u_v = his_vector.repeat(1,i_ids.shape[1]).view(i_ids.shape[0],i_ids.shape[1],-1) 84 | i_v = i_vectors 85 | 86 | return {'prediction': prediction.view(batch_size, -1), 'u_v': u_v, 'i_v':i_v} 87 | 88 | 89 | class SASRec(SequentialModel, SASRecBase): 90 | reader = 'SeqReader' 91 | runner = 'BaseRunner' 92 | extra_log_args = ['emb_size', 'num_layers', 'num_heads'] 93 | 94 | @staticmethod 95 | def parse_model_args(parser): 96 | parser = SASRecBase.parse_model_args(parser) 97 | return SequentialModel.parse_model_args(parser) 98 | 99 | def __init__(self, args, corpus): 100 | SequentialModel.__init__(self, args, corpus) 101 | self._base_init(args, corpus) 102 | 103 | def forward(self, feed_dict): 104 | out_dict = SASRecBase.forward(self, feed_dict) 105 | return {'prediction': out_dict['prediction']} 106 | 107 | class SASRecImpression(ImpressionSeqModel, SASRecBase): 108 | reader = 'ImpressionSeqReader' 109 | runner = 'ImpressionRunner' 110 | extra_log_args = ['emb_size', 'num_layers', 'num_heads'] 111 | 112 | @staticmethod 113 | def parse_model_args(parser): 114 | parser = SASRecBase.parse_model_args(parser) 115 | return ImpressionSeqModel.parse_model_args(parser) 116 | 117 | def __init__(self, args, corpus): 118 | ImpressionSeqModel.__init__(self, args, corpus) 119 | self._base_init(args, corpus) 120 | 121 | def forward(self, feed_dict): 122 | return SASRecBase.forward(self, feed_dict) -------------------------------------------------------------------------------- /src/models/sequential/SLRCPlus.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | # @Author : Chenyang Wang 3 | # @Email : THUwangcy@gmail.com 4 | 5 | """ SLRC+ 6 | Reference: 7 | "Modeling Item-specific Temporal Dynamics of Repeat Consumption for Recommender Systems" 8 | Chenyang Wang et al., TheWebConf'2019. 9 | Reference code: 10 | The authors' tensorflow implementation https://github.com/THUwangcy/SLRC 11 | Note: 12 | We generalize the original SLRC by also including mutual-excitation of relational history interactions. 13 | This makes SLRC+ a knowledge-aware model, and the original SLRC can be seen that there is only one special 14 | relation between items and themselves (i.e., repeat consumption). 15 | CMD example: 16 | python main.py --model_name SLRCPlus --emb_size 64 --lr 5e-4 --l2 1e-5 --dataset 'Grocery_and_Gourmet_Food' 17 | """ 18 | 19 | import torch 20 | import torch.nn as nn 21 | import torch.distributions 22 | import numpy as np 23 | 24 | from models.BaseModel import SequentialModel 25 | from helpers.KGReader import KGReader 26 | 27 | 28 | class SLRCPlus(SequentialModel): 29 | reader = 'KGReader' 30 | runner = 'BaseRunner' 31 | extra_log_args = ['emb_size'] 32 | 33 | @staticmethod 34 | def parse_model_args(parser): 35 | parser.add_argument('--emb_size', type=int, default=64, 36 | help='Size of embedding vectors.') 37 | parser.add_argument('--time_scalar', type=int, default=60 * 60 * 24 * 100, 38 | help='Time scalar for time intervals.') 39 | return SequentialModel.parse_model_args(parser) 40 | 41 | def __init__(self, args, corpus: KGReader): 42 | super().__init__(args, corpus) 43 | self.emb_size = args.emb_size 44 | self.time_scalar = args.time_scalar 45 | self.relation_num = len(corpus.item_relations) + 1 46 | self._define_params() 47 | self.apply(self.init_weights) 48 | 49 | def _define_params(self): 50 | self.u_embeddings = nn.Embedding(self.user_num, self.emb_size) 51 | self.i_embeddings = nn.Embedding(self.item_num, self.emb_size) 52 | self.user_bias = nn.Embedding(self.user_num, 1) 53 | self.item_bias = nn.Embedding(self.item_num, 1) 54 | 55 | self.global_alpha = nn.Parameter(torch.tensor(0.)) 56 | self.alphas = nn.Embedding(self.item_num, self.relation_num) 57 | self.pis = nn.Embedding(self.item_num, self.relation_num) 58 | self.betas = nn.Embedding(self.item_num, self.relation_num) 59 | self.sigmas = nn.Embedding(self.item_num, self.relation_num) 60 | self.mus = nn.Embedding(self.item_num, self.relation_num) 61 | 62 | def forward(self, feed_dict): 63 | self.check_list = [] 64 | u_ids = feed_dict['user_id'] # [batch_size] 65 | i_ids = feed_dict['item_id'] # [batch_size, -1] 66 | r_intervals = feed_dict['relational_interval'] # [batch_size, -1, relation_num] 67 | 68 | # Excitation 69 | alphas = self.global_alpha + self.alphas(i_ids) 70 | pis, mus = self.pis(i_ids) + 0.5, self.mus(i_ids) + 1 71 | betas = (self.betas(i_ids) + 1).clamp(min=1e-10, max=10) 72 | sigmas = (self.sigmas(i_ids) + 1).clamp(min=1e-10, max=10) 73 | mask = (r_intervals >= 0).float() 74 | delta_t = r_intervals * mask 75 | norm_dist = torch.distributions.normal.Normal(mus, sigmas) 76 | exp_dist = torch.distributions.exponential.Exponential(betas, validate_args=False) 77 | decay = pis * exp_dist.log_prob(delta_t).exp() + (1 - pis) * norm_dist.log_prob(delta_t).exp() 78 | excitation = (alphas * decay * mask).sum(-1) # [batch_size, -1] 79 | 80 | # Base Intensity (MF) 81 | u_bias = self.user_bias(u_ids) 82 | i_bias = self.item_bias(i_ids).squeeze(-1) 83 | cf_u_vectors = self.u_embeddings(u_ids) 84 | cf_i_vectors = self.i_embeddings(i_ids) 85 | base_intensity = (cf_u_vectors[:, None, :] * cf_i_vectors).sum(-1) 86 | base_intensity = base_intensity + u_bias + i_bias 87 | 88 | prediction = base_intensity + excitation 89 | return {'prediction': prediction.view(feed_dict['batch_size'], -1)} 90 | 91 | class Dataset(SequentialModel.Dataset): 92 | def _get_feed_dict(self, index): 93 | feed_dict = super()._get_feed_dict(index) 94 | user_id, time = self.data['user_id'][index], self.data['time'][index] 95 | history_item, history_time = feed_dict['history_items'], feed_dict['history_times'] 96 | 97 | # Collect time information related to the target item: 98 | # - re-consuming time gaps 99 | # - time intervals w.r.t. recent relational interactions 100 | relational_interval = list() 101 | for i, target_item in enumerate(feed_dict['item_id']): 102 | interval = np.ones(self.model.relation_num, dtype=float) * -1 # -1 if not existing 103 | # the first dimension for re-consuming time gaps 104 | for j in range(len(history_item))[::-1]: 105 | if history_item[j] == target_item: 106 | interval[0] = (time - history_time[j]) / self.model.time_scalar 107 | break 108 | # the rest for relational time intervals 109 | for r_idx in range(1, self.model.relation_num): 110 | for j in range(len(history_item))[::-1]: 111 | if (history_item[j], r_idx, target_item) in self.corpus.triplet_set: 112 | interval[r_idx] = (time - history_time[j]) / self.model.time_scalar 113 | break 114 | relational_interval.append(interval) 115 | feed_dict['relational_interval'] = np.array(relational_interval, dtype=np.float32) 116 | return feed_dict 117 | -------------------------------------------------------------------------------- /src/models/sequential/__init__.py: -------------------------------------------------------------------------------- 1 | from os.path import dirname, basename, isfile, join 2 | import glob 3 | 4 | modules = glob.glob(join(dirname(__file__), "*.py")) 5 | __all__ = [ 6 | basename(f)[:-3] for f in modules if isfile(f) and not f.endswith('__init__.py') 7 | ] 8 | -------------------------------------------------------------------------------- /src/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/THUwangcy/ReChorus/e58d752558634d42be60de561d6f74df748583cc/src/utils/__init__.py -------------------------------------------------------------------------------- /src/utils/utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: UTF-8 -*- 2 | 3 | import os 4 | import random 5 | import logging 6 | import torch 7 | import datetime 8 | import numpy as np 9 | import pandas as pd 10 | from typing import List, Dict, NoReturn, Any 11 | 12 | 13 | def init_seed(seed): 14 | random.seed(seed) 15 | np.random.seed(seed) 16 | torch.manual_seed(seed) 17 | torch.cuda.manual_seed(seed) 18 | torch.cuda.manual_seed_all(seed) 19 | torch.backends.cudnn.benchmark = False 20 | torch.backends.cudnn.deterministic = True 21 | 22 | 23 | def df_to_dict(df: pd.DataFrame) -> dict: 24 | res = df.to_dict('list') 25 | for key in res: 26 | res[key] = np.array(res[key]) 27 | return res 28 | 29 | 30 | def batch_to_gpu(batch: dict, device) -> dict: 31 | for c in batch: 32 | if type(batch[c]) is torch.Tensor: 33 | batch[c] = batch[c].to(device) 34 | return batch 35 | 36 | 37 | def check(check_list: List[tuple]) -> NoReturn: 38 | # observe selected tensors during training. 39 | logging.info('') 40 | for i, t in enumerate(check_list): 41 | d = np.array(t[1].detach().cpu()) 42 | logging.info(os.linesep.join( 43 | [t[0] + '\t' + str(d.shape), np.array2string(d, threshold=20)] 44 | ) + os.linesep) 45 | 46 | 47 | def eval_list_columns(df: pd.DataFrame) -> pd.DataFrame: 48 | for col in df.columns: 49 | if pd.api.types.is_string_dtype(df[col]): 50 | df[col] = df[col].apply(lambda x: eval(str(x))) # some list-value columns 51 | return df 52 | 53 | 54 | def format_metric(result_dict: Dict[str, Any]) -> str: 55 | assert type(result_dict) == dict 56 | format_str = [] 57 | metrics = np.unique([k.split('@')[0] for k in result_dict.keys()]) 58 | topks = np.unique([int(k.split('@')[1]) for k in result_dict.keys() if '@' in k]) 59 | if not len(topks): 60 | topks = ['All'] 61 | for topk in np.sort(topks): 62 | for metric in np.sort(metrics): 63 | name = '{}@{}'.format(metric, topk) 64 | m = result_dict[name] if topk != 'All' else result_dict[metric] 65 | if type(m) is float or type(m) is np.float or type(m) is np.float32 or type(m) is np.float64: 66 | format_str.append('{}:{:<.4f}'.format(name, m)) 67 | elif type(m) is int or type(m) is np.int or type(m) is np.int32 or type(m) is np.int64: 68 | format_str.append('{}:{}'.format(name, m)) 69 | return ','.join(format_str) 70 | 71 | 72 | def format_arg_str(args, exclude_lst: list, max_len=20) -> str: 73 | linesep = os.linesep 74 | arg_dict = vars(args) 75 | keys = [k for k in arg_dict.keys() if k not in exclude_lst] 76 | values = [arg_dict[k] for k in keys] 77 | key_title, value_title = 'Arguments', 'Values' 78 | key_max_len = max(map(lambda x: len(str(x)), keys)) 79 | value_max_len = min(max(map(lambda x: len(str(x)), values)), max_len) 80 | key_max_len, value_max_len = max([len(key_title), key_max_len]), max([len(value_title), value_max_len]) 81 | horizon_len = key_max_len + value_max_len + 5 82 | res_str = linesep + '=' * horizon_len + linesep 83 | res_str += ' ' + key_title + ' ' * (key_max_len - len(key_title)) + ' | ' \ 84 | + value_title + ' ' * (value_max_len - len(value_title)) + ' ' + linesep + '=' * horizon_len + linesep 85 | for key in sorted(keys): 86 | value = arg_dict[key] 87 | if value is not None: 88 | key, value = str(key), str(value).replace('\t', '\\t') 89 | value = value[:max_len-3] + '...' if len(value) > max_len else value 90 | res_str += ' ' + key + ' ' * (key_max_len - len(key)) + ' | ' \ 91 | + value + ' ' * (value_max_len - len(value)) + linesep 92 | res_str += '=' * horizon_len 93 | return res_str 94 | 95 | 96 | def check_dir(file_name: str): 97 | dir_path = os.path.dirname(file_name) 98 | if not os.path.exists(dir_path): 99 | print('make dirs:', dir_path) 100 | os.makedirs(dir_path) 101 | 102 | 103 | def non_increasing(lst: list) -> bool: 104 | return all(x >= y for x, y in zip([lst[0]]*(len(lst)-1), lst[1:])) # update the calculation of non_increasing to fit ealry stopping, 2023.5.14, Jiayu Li 105 | 106 | 107 | def get_time(): 108 | return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S") 109 | 110 | --------------------------------------------------------------------------------