Add files using upload-large-folder tool

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+asset/model.jpg filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,124 @@
+# LC-Rec
+
+This is the official PyTorch implementation for the paper:
+
+> [Adapting Large Language Models by Integrating Collaborative Semantics for Recommendation](https://arxiv.org/abs/2311.09049)
+
+## Overview
+
+We propose **LC-Rec**, a new approach to integrate **L**anguage and **C**ollaborative semantics for improving LLMs in **Rec**ommender systems. To tackle the large gap between the language semantics modeled by LLMs and collaborative semantics implied by recommender systems, we make two major contributions in two aspects. For item indexing, we design a learning-based vector quantization method with uniform semantic mapping, which can assign meaningful and non-conflicting IDs (called item indices) for items. For alignment tuning, we propose a series of specially designed tuning tasks to enhance the integration of collaborative semantics in LLMs. Our fine-tuning tasks enforce LLMs to deeply integrate language and collaborative semantics (characterized by the learned item indices), so as to achieve an effective adaptation to recommender systems.
+
+![model](./asset/model.jpg)
+
+## Requirements
+
+```
+torch==1.13.1+cu117
+accelerate
+bitsandbytes
+deepspeed
+evaluate
+peft
+sentencepiece
+tqdm
+transformers
+```
+
+## Model Checkpoint
+
+The delta weights on the three datasets can be downloaded from huggingface hub ([Instruments](https://huggingface.co/bwzheng0324/lc-rec-instruments-delta), [Arts](https://huggingface.co/bwzheng0324/lc-rec-arts-delta), [Games](https://huggingface.co/bwzheng0324/lc-rec-games-delta)). After downloading, you can add our deltas to the original LLaMA weights to obtain LC-Rec weights:
+
+1. Get the original [LLaMA](https://huggingface.co/huggyllama/llama-7b) weights.
+2. Use the following scripts to get LC-Rec weights by applying our delta.
+
+```shell
+python -m convert/merge_delta.py \
+    --base-model-path /path/to/llama-7b \
+    --target-model-path /path/output/lc-rec \
+    --delta-path bwzheng0324/lc-rec-games-delta
+```
+
+## Dataset
+
+We use three datasets in our paper, all of which have been uploaded to [Google Drive](https://drive.google.com/drive/folders/1RcJ2M1l5zWPHYuGd9l5Gibcs5w5aI3y6?usp=sharing) 
+
+## Train
+
+The detailed scripts for all three datasets are in `run.sh`:
+
+```shell
+DATASET=Games
+BASE_MODEL=huggyllama/llama-7b
+DATA_PATH=./data
+OUTPUT_DIR=./ckpt/$DATASET/
+
+torchrun --nproc_per_node=8 --master_port=3324 finetune.py \
+    --base_model $BASE_MODEL \
+    --output_dir $OUTPUT_DIR \
+    --dataset $DATASET \
+    --data_path $DATA_PATH \
+    --per_device_batch_size 8 \
+    --gradient_accumulation_steps 2 \
+    --learning_rate 5e-5 \
+    --epochs 4 \
+    --weight_decay 0.01 \
+    --save_and_eval_strategy epoch \
+    --deepspeed ./config/ds_z3_bf16.json \
+    --bf16 \
+    --only_train_response \
+    --tasks seqrec,item2index,index2item,fusionseqrec,itemsearch,preferenceobtain \
+    --train_prompt_sample_num 1,1,1,1,1,1 \
+    --train_data_sample_num 0,0,0,100000,0,0 \
+    --index_file .index.json
+
+
+cd convert
+nohup ./convert.sh $OUTPUT_DIR >convert.log 2>&1 &
+cd ..
+```
+
+## Test
+
+Test with a single GPU:
+
+```shell
+DATASET=Games
+DATA_PATH=./data
+OUTPUT_DIR=./ckpt/$DATASET/
+RESULTS_FILE=./results/$DATASET/result.json
+
+python test.py \
+    --gpu_id 0 \
+    --ckpt_path $CKPT_PATH \
+    --dataset $DATASET \
+    --data_path $DATA_PATH \
+    --results_file $RESULTS_FILE \
+    --test_batch_size 1 \
+    --num_beams 20 \
+    --test_prompt_ids all \
+    --index_file .index.json
+```
+
+Test with multiple GPUs:
+
+```shell
+DATASET=Games
+DATA_PATH=./data
+OUTPUT_DIR=./ckpt/$DATASET/
+RESULTS_FILE=./results/$DATASET/result.json
+
+torchrun --nproc_per_node=8 --master_port=4324 test_ddp.py \
+    --ckpt_path $CKPT_PATH \
+    --dataset $DATASET \
+    --data_path $DATA_PATH \
+    --results_file $RESULTS_FILE \
+    --test_batch_size 1 \
+    --num_beams 20 \
+    --test_prompt_ids all \
+    --index_file .index.json
+```
+
+## Acknowledgement
+
+The implementation is based on [HuggingFace](https://github.com/huggingface/transformers).
+

collator.py

@@ -0,0 +1,77 @@
+import torch
+import copy
+import argparse
+from dataclasses import dataclass
+
+import transformers
+import math
+from torch.utils.data import Sampler
+import torch.distributed as dist
+from transformers import LlamaForCausalLM, LlamaTokenizer, LlamaConfig, T5Tokenizer, T5Config, T5ForConditionalGeneration
+
+
+class Collator(object):
+
+    def __init__(self, args, tokenizer):
+        self.args = args
+        self.only_train_response = args.only_train_response
+        self.tokenizer = tokenizer
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token_id = self.tokenizer.unk_token_id
+        # print(self.tokenizer.model_max_length)
+
+    def __call__(self, batch):
+
+        input_texts = [d["input_ids"] for d in batch]
+        full_texts = [d["labels"] + self.tokenizer.eos_token for d in batch]
+
+        inputs = self.tokenizer(
+            text = full_texts,
+            text_target = input_texts,
+            return_tensors="pt",
+            padding="longest",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_attention_mask=True,
+        )
+        labels = copy.deepcopy(inputs["input_ids"])
+        if self.only_train_response:
+            # ignore padding
+            labels[labels == self.tokenizer.pad_token_id] = -100
+            # ignore input text
+            labels[torch.where(inputs["labels"] != self.tokenizer.pad_token_id)] = -100
+
+        inputs["labels"] = labels
+
+
+        return inputs
+
+
+
+class TestCollator(object):
+
+    def __init__(self, args, tokenizer):
+        self.args = args
+        self.tokenizer = tokenizer
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token_id = 0
+
+        if isinstance(self.tokenizer, LlamaTokenizer):
+            # Allow batched inference
+            self.tokenizer.padding_side = "left"
+
+    def __call__(self, batch):
+
+        input_texts = [d["input_ids"] for d in batch]
+        targets = [d["labels"] for d in batch]
+        inputs = self.tokenizer(
+            text=input_texts,
+            return_tensors="pt",
+            padding="longest",
+            max_length=self.tokenizer.model_max_length,
+            truncation=True,
+            return_attention_mask=True,
+        )
+
+        return (inputs, targets)
+

config/ds_z2_bf16.json

@@ -0,0 +1,28 @@
+{
+    "bf16": {
+      "enabled": "auto"
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 5e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 5e8,
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false,
+    "flops_profiler": {
+        "enabled": true,
+        "profile_step": 10,
+        "module_depth": -1,
+        "top_modules": 3,
+        "detailed": true,
+        "output_file": "flops_profiler.out"
+    }
+}

config/ds_z2_fp16.json

@@ -0,0 +1,34 @@
+{
+    "fp16": {
+      "enabled": "auto",
+      "auto_cast": false,
+      "loss_scale": 0,
+      "initial_scale_power": 16,
+      "loss_scale_window": 1000,
+      "hysteresis": 2,
+      "min_loss_scale": 1
+    },
+    "zero_optimization": {
+        "stage": 2,
+        "allgather_partitions": true,
+        "allgather_bucket_size": 5e8,
+        "overlap_comm": true,
+        "reduce_scatter": true,
+        "reduce_bucket_size": 5e8,
+        "contiguous_gradients": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false,
+    "flops_profiler": {
+        "enabled": true,
+        "profile_step": 10,
+        "module_depth": -1,
+        "top_modules": 3,
+        "detailed": true,
+        "output_file": "flops_profiler.out"
+    }
+}

config/ds_z3_bf16.json

@@ -0,0 +1,31 @@
+{
+    "bf16": {
+      "enabled": "auto"
+    },
+    "zero_optimization": {
+      "stage": 3,
+      "overlap_comm": true,
+      "contiguous_gradients": true,
+      "sub_group_size": 1e9,
+      "reduce_bucket_size": "auto",
+      "stage3_prefetch_bucket_size": "auto",
+      "stage3_param_persistence_threshold": "auto",
+      "stage3_max_live_parameters": 1e9,
+      "stage3_max_reuse_distance": 1e9,
+      "stage3_gather_16bit_weights_on_model_save": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false,
+    "flops_profiler": {
+        "enabled": true,
+        "profile_step": 10,
+        "module_depth": -1,
+        "top_modules": 3,
+        "detailed": true,
+        "output_file": "flops_profiler.out"
+    }
+}

config/ds_z3_bf16_save16bit.json

@@ -0,0 +1,31 @@
+{
+    "bf16": {
+      "enabled": "auto"
+    },
+    "zero_optimization": {
+      "stage": 3,
+      "overlap_comm": true,
+      "contiguous_gradients": true,
+      "sub_group_size": 1e9,
+      "reduce_bucket_size": "auto",
+      "stage3_prefetch_bucket_size": "auto",
+      "stage3_param_persistence_threshold": "auto",
+      "stage3_max_live_parameters": 1e9,
+      "stage3_max_reuse_distance": 1e9,
+      "stage3_gather_16bit_weights_on_model_save": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false,
+    "flops_profiler": {
+        "enabled": true,
+        "profile_step": 10,
+        "module_depth": -1,
+        "top_modules": 3,
+        "detailed": true,
+        "output_file": "flops_profiler.out"
+    }
+}

config/ds_z3_fp16.json

@@ -0,0 +1,37 @@
+{
+    "fp16": {
+      "enabled": "auto",
+      "auto_cast": false,
+      "loss_scale": 0,
+      "initial_scale_power": 16,
+      "loss_scale_window": 1000,
+      "hysteresis": 2,
+      "min_loss_scale": 1
+    },
+    "zero_optimization": {
+      "stage": 3,
+      "overlap_comm": true,
+      "contiguous_gradients": true,
+      "sub_group_size": 1e9,
+      "reduce_bucket_size": "auto",
+      "stage3_prefetch_bucket_size": "auto",
+      "stage3_param_persistence_threshold": "auto",
+      "stage3_max_live_parameters": 1e9,
+      "stage3_max_reuse_distance": 1e9,
+      "stage3_gather_16bit_weights_on_model_save": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false,
+    "flops_profiler": {
+        "enabled": true,
+        "profile_step": 10,
+        "module_depth": -1,
+        "top_modules": 3,
+        "detailed": true,
+        "output_file": "flops_profiler.out"
+    }
+}

config/ds_z3_fp16_save16bit.json

@@ -0,0 +1,37 @@
+{
+    "fp16": {
+      "enabled": "auto",
+      "auto_cast": false,
+      "loss_scale": 0,
+      "initial_scale_power": 16,
+      "loss_scale_window": 1000,
+      "hysteresis": 2,
+      "min_loss_scale": 1
+    },
+    "zero_optimization": {
+      "stage": 3,
+      "overlap_comm": true,
+      "contiguous_gradients": true,
+      "sub_group_size": 1e9,
+      "reduce_bucket_size": "auto",
+      "stage3_prefetch_bucket_size": "auto",
+      "stage3_param_persistence_threshold": "auto",
+      "stage3_max_live_parameters": 1e9,
+      "stage3_max_reuse_distance": 1e9,
+      "stage3_gather_16bit_weights_on_model_save": true
+    },
+    "gradient_accumulation_steps": "auto",
+    "gradient_clipping": "auto",
+    "steps_per_print": 2000,
+    "train_batch_size": "auto",
+    "train_micro_batch_size_per_gpu": "auto",
+    "wall_clock_breakdown": false,
+    "flops_profiler": {
+        "enabled": true,
+        "profile_step": 10,
+        "module_depth": -1,
+        "top_modules": 3,
+        "detailed": true,
+        "output_file": "flops_profiler.out"
+    }
+}

convert/convert.py

@@ -0,0 +1,16 @@
+import transformers
+import argparse
+import os
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--source", "-s", type=str, default="", help="source path of models")
+    parser.add_argument("--target", "-t", type=str, default="", help="target path of models")
+
+    args, _ = parser.parse_known_args()
+
+    assert os.path.exists(args.source)
+    assert args.target != ""
+
+    model = transformers.AutoModelForCausalLM.from_pretrained(args.source)
+    model.save_pretrained(args.target, state_dict=model.state_dict())

convert/convert.sh

@@ -0,0 +1,18 @@
+model=$1
+
+set -x
+
+for step in `ls ${model} | grep checkpoint | awk -F'-' '{ print $2 }'`
+do
+mkdir ${model}/tmp-checkpoint-${step}
+mkdir ${model}/final-checkpoint-${step}
+python ./zero_to_fp32.py ${model}/checkpoint-${step}/ ${model}/tmp-checkpoint-${step}/pytorch_model.bin
+cp ${model}/*.json ${model}/tmp-checkpoint-${step}
+python ./convert.py -s ${model}/tmp-checkpoint-${step} -t ${model}/final-checkpoint-${step}
+cp ${model}/checkpoint-${step}/*.json ${model}/final-checkpoint-${step}
+cp ${model}/*.json ${model}/final-checkpoint-${step}
+cp ${model}/tokenizer* ${model}/final-checkpoint-${step}
+cp ${model}/train* ${model}/final-checkpoint-${step}
+#rm -rf ${model}/tmp-checkpoint-${step} ${model}/checkpoint-${step} ${model}/global_step${step}
+#mv ${model}/final-checkpoint-${step} ${model}/checkpoint-${step}
+done

convert/convert_fp16.py

@@ -0,0 +1,23 @@
+
+import argparse
+
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+
+def convert_fp16(in_checkpoint, out_checkpoint):
+    tokenizer = AutoTokenizer.from_pretrained(in_checkpoint, use_fast=False)
+    model = AutoModelForCausalLM.from_pretrained(
+        in_checkpoint, torch_dtype=torch.float16, low_cpu_mem_usage=True
+    )
+    model.save_pretrained(out_checkpoint)
+    tokenizer.save_pretrained(out_checkpoint)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--in-checkpoint", type=str, help="Path to the model")
+    parser.add_argument("--out-checkpoint", type=str, help="Path to the output model")
+    args = parser.parse_args()
+
+    convert_fp16(args.in_checkpoint, args.out_checkpoint)

convert/make_delta.py

@@ -0,0 +1,46 @@
+
+import argparse
+
+import torch
+from tqdm import tqdm
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+
+def make_delta(base_model_path, target_model_path, delta_path):
+    print(f"Loading the base model from {base_model_path}")
+    base = AutoModelForCausalLM.from_pretrained(
+        base_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True
+    )
+
+    print(f"Loading the target model from {target_model_path}")
+    target = AutoModelForCausalLM.from_pretrained(
+        target_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True
+    )
+    target_tokenizer = AutoTokenizer.from_pretrained(target_model_path, use_fast=False)
+
+    print("Calculating the delta")
+    for name, param in tqdm(target.state_dict().items(), desc="Calculating delta"):
+        assert name in base.state_dict()
+        if param.shape == base.state_dict()[name].shape:
+            param.data -= base.state_dict()[name]
+        else:
+            print(name)
+
+    print(f"Saving the delta to {delta_path}")
+    if args.hub_repo_id:
+        kwargs = {"push_to_hub": True, "repo_id": args.hub_repo_id}
+    else:
+        kwargs = {}
+    target.save_pretrained(delta_path, **kwargs)
+    target_tokenizer.save_pretrained(delta_path, **kwargs)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--base-model-path", type=str, required=True)
+    parser.add_argument("--target-model-path", type=str, required=True)
+    parser.add_argument("--delta-path", type=str, required=True)
+    parser.add_argument("--hub-repo-id", type=str)
+    args = parser.parse_args()
+
+    make_delta(args.base_model_path, args.target_model_path, args.delta_path)

convert/merge_delta.py

@@ -0,0 +1,167 @@
+
+import argparse
+import gc
+import glob
+import json
+import os
+import shutil
+import tempfile
+
+from huggingface_hub import snapshot_download
+import torch
+from torch import nn
+from tqdm import tqdm
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
+
+
+GB = 1 << 30
+
+
+def split_files(model_path, tmp_path, split_size):
+    if not os.path.exists(model_path):
+        model_path = snapshot_download(repo_id=model_path)
+    if not os.path.exists(tmp_path):
+        os.makedirs(tmp_path)
+
+    file_pattern = os.path.join(model_path, "pytorch_model-*.bin")
+    files = glob.glob(file_pattern)
+
+    part = 0
+    try:
+        for file_path in tqdm(files):
+            state_dict = torch.load(file_path)
+            new_state_dict = {}
+
+            current_size = 0
+            for name, param in state_dict.items():
+                param_size = param.numel() * param.element_size()
+
+                if current_size + param_size > split_size:
+                    new_file_name = f"pytorch_model-{part}.bin"
+                    new_file_path = os.path.join(tmp_path, new_file_name)
+                    torch.save(new_state_dict, new_file_path)
+                    current_size = 0
+                    new_state_dict = None
+                    gc.collect()
+                    new_state_dict = {}
+                    part += 1
+
+                new_state_dict[name] = param
+                current_size += param_size
+
+            new_file_name = f"pytorch_model-{part}.bin"
+            new_file_path = os.path.join(tmp_path, new_file_name)
+            torch.save(new_state_dict, new_file_path)
+            new_state_dict = None
+            gc.collect()
+            new_state_dict = {}
+            part += 1
+    except Exception as e:
+        print(f"An error occurred during split_files: {e}")
+        shutil.rmtree(tmp_path)
+        raise
+
+
+def apply_delta_low_cpu_mem(base_model_path, target_model_path, delta_path):
+    delta_tokenizer = AutoTokenizer.from_pretrained(delta_path, use_fast=False)
+    delta_config = AutoConfig.from_pretrained(delta_path)
+
+    if os.path.exists(target_model_path):
+        shutil.rmtree(target_model_path)
+    os.makedirs(target_model_path)
+
+    split_size = 4 * GB
+
+    with tempfile.TemporaryDirectory() as tmp_base_path, tempfile.TemporaryDirectory() as tmp_delta_path:
+        print(f"Split files for the base model to {tmp_base_path}")
+        split_files(base_model_path, tmp_base_path, split_size)
+        print(f"Split files for the delta weights to {tmp_delta_path}")
+        split_files(delta_path, tmp_delta_path, split_size)
+
+        base_pattern = os.path.join(tmp_base_path, "pytorch_model-*.bin")
+        base_files = glob.glob(base_pattern)
+        base_state_dict = torch.load(base_files[0])
+        delta_pattern = os.path.join(tmp_delta_path, "pytorch_model-*.bin")
+        delta_files = glob.glob(delta_pattern)
+        # delta_state_dict = torch.load(delta_files[0])
+
+        print("Applying the delta")
+        weight_map = {}
+        total_size = 0
+
+        for i, delta_file in tqdm(enumerate(delta_files)):
+            state_dict = torch.load(delta_file)
+            file_name = f"pytorch_model-{i}.bin"
+            for name, param in state_dict.items():
+                if name not in base_state_dict:
+                    for base_file in base_files:
+                        base_state_dict = torch.load(base_file)
+                        gc.collect()
+                        if name in base_state_dict:
+                            break
+                if state_dict[name].shape == base_state_dict[name].shape:
+                    state_dict[name] += base_state_dict[name]
+                else:
+                    print(name)
+                weight_map[name] = file_name
+                total_size += param.numel() * param.element_size()
+                gc.collect()
+            torch.save(state_dict, os.path.join(target_model_path, file_name))
+
+        with open(
+            os.path.join(target_model_path, "pytorch_model.bin.index.json"), "w"
+        ) as f:
+            json.dump(
+                {"weight_map": weight_map, "metadata": {"total_size": total_size}}, f
+            )
+
+    print(f"Saving the target model to {target_model_path}")
+    delta_tokenizer.save_pretrained(target_model_path)
+    delta_config.save_pretrained(target_model_path)
+
+
+def apply_delta(base_model_path, target_model_path, delta_path):
+    print(f"Loading the delta weights from {delta_path}")
+    delta_tokenizer = AutoTokenizer.from_pretrained(delta_path, use_fast=False)
+    delta = AutoModelForCausalLM.from_pretrained(
+        delta_path, torch_dtype=torch.float16, low_cpu_mem_usage=True
+    )
+
+    print(f"Loading the base model from {base_model_path}")
+    base = AutoModelForCausalLM.from_pretrained(
+        base_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True
+    )
+
+    print("Applying the delta")
+    for name, param in tqdm(delta.state_dict().items(), desc="Applying delta"):
+        assert name in base.state_dict()
+        if param.shape == base.state_dict()[name].shape:
+            param.data += base.state_dict()[name]
+        else:
+            print(name)
+
+
+    print(f"Saving the target model to {target_model_path}")
+    delta.save_pretrained(target_model_path)
+    delta_tokenizer.save_pretrained(target_model_path)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--base-model-path", type=str, required=True)
+    parser.add_argument("--target-model-path", type=str, required=True)
+    parser.add_argument("--delta-path", type=str, required=True)
+    parser.add_argument(
+        "--low-cpu-mem",
+        action="store_true",
+        help="Lower the cpu memory usage. This will split large files and use "
+        "disk as swap to reduce the memory usage below 10GB.",
+    )
+    args = parser.parse_args()
+
+    if args.low_cpu_mem:
+        apply_delta_low_cpu_mem(
+            args.base_model_path, args.target_model_path, args.delta_path
+        )
+    else:
+        apply_delta(args.base_model_path, args.target_model_path, args.delta_path)

convert/zero_to_fp32.py

@@ -0,0 +1,600 @@
+#!/usr/bin/env python
+
+# Copyright (c) Microsoft Corporation.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+# This script extracts fp32 consolidated weights from a zero 2 and 3 DeepSpeed checkpoints. It gets
+# copied into the top level checkpoint dir, so the user can easily do the conversion at any point in
+# the future. Once extracted, the weights don't require DeepSpeed and can be used in any
+# application.
+#
+# example: python zero_to_fp32.py . pytorch_model.bin
+
+import argparse
+import torch
+import glob
+import math
+import os
+import re
+from collections import OrderedDict
+from dataclasses import dataclass
+from tqdm import tqdm
+
+# while this script doesn't use deepspeed to recover data, since the checkpoints are pickled with
+# DeepSpeed data structures it has to be available in the current python environment.
+from deepspeed.utils import logger
+from deepspeed.checkpoint.constants import (DS_VERSION, OPTIMIZER_STATE_DICT, SINGLE_PARTITION_OF_FP32_GROUPS,
+                                            FP32_FLAT_GROUPS, ZERO_STAGE, PARTITION_COUNT, PARAM_SHAPES, BUFFER_NAMES,
+                                            FROZEN_PARAM_SHAPES, FROZEN_PARAM_FRAGMENTS)
+
+
+@dataclass
+class zero_model_state:
+    buffers: dict()
+    param_shapes: dict()
+    shared_params: list
+    ds_version: int
+    frozen_param_shapes: dict()
+    frozen_param_fragments: dict()
+
+
+debug = 0
+
+# load to cpu
+device = torch.device('cpu')
+
+
+def atoi(text):
+    return int(text) if text.isdigit() else text
+
+
+def natural_keys(text):
+    '''
+    alist.sort(key=natural_keys) sorts in human order
+    http://nedbatchelder.com/blog/200712/human_sorting.html
+    (See Toothy's implementation in the comments)
+    '''
+    return [atoi(c) for c in re.split(r'(\d+)', text)]
+
+
+def get_model_state_file(checkpoint_dir, zero_stage):
+    if not os.path.isdir(checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{checkpoint_dir}' doesn't exist")
+
+    # there should be only one file
+    if zero_stage == 2:
+        file = os.path.join(checkpoint_dir, "mp_rank_00_model_states.pt")
+    elif zero_stage == 3:
+        file = os.path.join(checkpoint_dir, "zero_pp_rank_0_mp_rank_00_model_states.pt")
+
+    if not os.path.exists(file):
+        raise FileNotFoundError(f"can't find model states file at '{file}'")
+
+    return file
+
+
+def get_checkpoint_files(checkpoint_dir, glob_pattern):
+    # XXX: need to test that this simple glob rule works for multi-node setup too
+    ckpt_files = sorted(glob.glob(os.path.join(checkpoint_dir, glob_pattern)), key=natural_keys)
+
+    if len(ckpt_files) == 0:
+        raise FileNotFoundError(f"can't find {glob_pattern} files in directory '{checkpoint_dir}'")
+
+    return ckpt_files
+
+
+def get_optim_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_optim_states.pt")
+
+
+def get_model_state_files(checkpoint_dir):
+    return get_checkpoint_files(checkpoint_dir, "*_model_states.pt")
+
+
+def parse_model_states(files):
+    zero_model_states = []
+    for file in files:
+        state_dict = torch.load(file, map_location=device)
+
+        if BUFFER_NAMES not in state_dict:
+            raise ValueError(f"{file} is not a model state checkpoint")
+        buffer_names = state_dict[BUFFER_NAMES]
+        if debug:
+            print("Found buffers:", buffer_names)
+
+        # recover just the buffers while restoring them to fp32 if they were saved in fp16
+        buffers = {k: v.float() for k, v in state_dict["module"].items() if k in buffer_names}
+        param_shapes = state_dict[PARAM_SHAPES]
+
+        # collect parameters that are included in param_shapes
+        param_names = []
+        for s in param_shapes:
+            for name in s.keys():
+                param_names.append(name)
+
+        # update with frozen parameters
+        frozen_param_shapes = state_dict.get(FROZEN_PARAM_SHAPES, None)
+        if frozen_param_shapes is not None:
+            if debug:
+                print(f"Found frozen_param_shapes: {frozen_param_shapes}")
+            param_names += list(frozen_param_shapes.keys())
+
+        # record shared parameters so that they can be recovered based on partners
+        # this is because such parameters holding reference only are not saved by optimizer
+        shared_params = []
+        for param in state_dict["module"]:
+            if param not in [*param_names, *buffer_names]:
+                for share_param in state_dict["module"]:
+                    if (state_dict["module"][share_param].data_ptr() == state_dict["module"][param].data_ptr()
+                            and share_param != param):
+                        shared_params.append([param, share_param])
+                        break
+
+        ds_version = state_dict.get(DS_VERSION, None)
+
+        frozen_param_fragments = state_dict.get(FROZEN_PARAM_FRAGMENTS, None)
+
+        z_model_state = zero_model_state(buffers=buffers,
+                                         param_shapes=param_shapes,
+                                         shared_params=shared_params,
+                                         ds_version=ds_version,
+                                         frozen_param_shapes=frozen_param_shapes,
+                                         frozen_param_fragments=frozen_param_fragments)
+        zero_model_states.append(z_model_state)
+
+    return zero_model_states
+
+
+def parse_optim_states(files, ds_checkpoint_dir):
+
+    total_files = len(files)
+    state_dicts = []
+    for i, f in enumerate(tqdm(files)):
+        state_dicts.append(torch.load(f, map_location=device))
+        if i == 0:
+            if not ZERO_STAGE in state_dicts[0][OPTIMIZER_STATE_DICT]:
+                raise ValueError(f"{files[0]} is not a zero checkpoint")
+            zero_stage = state_dicts[0][OPTIMIZER_STATE_DICT][ZERO_STAGE]
+            world_size = state_dicts[0][OPTIMIZER_STATE_DICT][PARTITION_COUNT]
+
+            # For ZeRO-2 each param group can have different partition_count as data parallelism for expert
+            # parameters can be different from data parallelism for non-expert parameters. So we can just
+            # use the max of the partition_count to get the dp world_size.
+
+            if type(world_size) is list:
+                world_size = max(world_size)
+
+            if world_size != total_files:
+                raise ValueError(
+                    f"Expected {world_size} of '*_optim_states.pt' under '{ds_checkpoint_dir}' but found {total_files} files. "
+                    "Possibly due to an overwrite of an old checkpoint, or a checkpoint didn't get saved by one or more processes."
+                )
+
+            # the groups are named differently in each stage
+            if zero_stage == 2:
+                fp32_groups_key = SINGLE_PARTITION_OF_FP32_GROUPS
+            elif zero_stage == 3:
+                fp32_groups_key = FP32_FLAT_GROUPS
+            else:
+                raise ValueError(f"unknown zero stage {zero_stage}")
+
+        key_list = list(state_dicts[-1][OPTIMIZER_STATE_DICT].keys())
+        for key in key_list:
+            if zero_stage == 2:
+                if key != fp32_groups_key:
+                    del state_dicts[-1][OPTIMIZER_STATE_DICT][key]
+            elif zero_stage == 3:
+                if key == fp32_groups_key:
+                    value = torch.cat(state_dicts[-1][OPTIMIZER_STATE_DICT][fp32_groups_key], 0)
+                del state_dicts[-1][OPTIMIZER_STATE_DICT][key]
+                if key == fp32_groups_key:
+                    state_dicts[-1][OPTIMIZER_STATE_DICT][key] = value
+
+    print('zero_stage:', zero_stage)
+    fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
+    # if zero_stage == 2:
+    #     # fp32_flat_groups = [state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key] for i in range(len(state_dicts))]
+    # elif zero_stage == 3:
+    #     # if there is more than one param group, there will be multiple flattened tensors - one
+    #     # flattened tensor per group - for simplicity merge them into a single tensor
+    #     #
+    #     # XXX: could make the script more memory efficient for when there are multiple groups - it
+    #     # will require matching the sub-lists of param_shapes for each param group flattened tensor
+
+    #     print('start!')
+    #     # fp32_flat_groups = [
+    #     #     torch.cat(state_dicts[i][OPTIMIZER_STATE_DICT][fp32_groups_key], 0) for i in range(len(state_dicts))
+    #     # ]
+
+    return zero_stage, world_size, fp32_flat_groups
+
+
+def _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir):
+    """
+    Returns fp32 state_dict reconstructed from ds checkpoint
+
+    Args:
+        - ``ds_checkpoint_dir``: path to the deepspeed checkpoint folder (where the optimizer files are)
+
+    """
+    print(f"Processing zero checkpoint '{ds_checkpoint_dir}'")
+
+    optim_files = get_optim_files(ds_checkpoint_dir)
+    zero_stage, world_size, fp32_flat_groups = parse_optim_states(optim_files, ds_checkpoint_dir)
+    print(f"Detected checkpoint of type zero stage {zero_stage}, world_size: {world_size}")
+
+    model_files = get_model_state_files(ds_checkpoint_dir)
+
+    zero_model_states = parse_model_states(model_files)
+    print(f'Parsing checkpoint created by deepspeed=={zero_model_states[0].ds_version}')
+
+    if zero_stage == 2:
+        return _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states)
+    elif zero_stage == 3:
+        return _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states)
+
+
+def _zero2_merge_frozen_params(state_dict, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+    frozen_param_fragments = zero_model_states[0].frozen_param_fragments
+
+    if debug:
+        num_elem = sum(s.numel() for s in frozen_param_shapes.values())
+        print(f'rank 0: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in frozen_param_fragments.values()])
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in frozen_param_shapes.items():
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        state_dict[name] = frozen_param_fragments[name]
+
+        if debug:
+            print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+
+    # Reconstruction protocol:
+    #
+    # XXX: document this
+
+    if debug:
+        for i in range(world_size):
+            for j in range(len(fp32_flat_groups[0])):
+                print(f"{FP32_FLAT_GROUPS}[{i}][{j}].shape={fp32_flat_groups[i][j].shape}")
+
+    # XXX: memory usage doubles here (zero2)
+    num_param_groups = len(fp32_flat_groups[0])
+    merged_single_partition_of_fp32_groups = []
+    for i in range(num_param_groups):
+        merged_partitions = [sd[i] for sd in fp32_flat_groups]
+        full_single_fp32_vector = torch.cat(merged_partitions, 0)
+        merged_single_partition_of_fp32_groups.append(full_single_fp32_vector)
+    avail_numel = sum(
+        [full_single_fp32_vector.numel() for full_single_fp32_vector in merged_single_partition_of_fp32_groups])
+
+    if debug:
+        wanted_params = sum([len(shapes) for shapes in param_shapes])
+        wanted_numel = sum([sum(shape.numel() for shape in shapes.values()) for shapes in param_shapes])
+        # not asserting if there is a mismatch due to possible padding
+        print(f"Have {avail_numel} numels to process.")
+        print(f"Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    total_numel = 0
+    total_params = 0
+    for shapes, full_single_fp32_vector in zip(param_shapes, merged_single_partition_of_fp32_groups):
+        offset = 0
+        avail_numel = full_single_fp32_vector.numel()
+        for name, shape in shapes.items():
+
+            unpartitioned_numel = shape.numel()
+            total_numel += unpartitioned_numel
+            total_params += 1
+
+            if debug:
+                print(f"{name} full shape: {shape} unpartitioned numel {unpartitioned_numel} ")
+            state_dict[name] = full_single_fp32_vector.narrow(0, offset, unpartitioned_numel).view(shape)
+            offset += unpartitioned_numel
+
+        # Z2 started to align to 2*world_size to improve nccl performance. Therefore both offset and
+        # avail_numel can differ by anywhere between 0..2*world_size. Due to two unrelated complex
+        # paddings performed in the code it's almost impossible to predict the exact numbers w/o the
+        # live optimizer object, so we are checking that the numbers are within the right range
+        align_to = 2 * world_size
+
+        def zero2_align(x):
+            return align_to * math.ceil(x / align_to)
+
+        if debug:
+            print(f"original offset={offset}, avail_numel={avail_numel}")
+
+        offset = zero2_align(offset)
+        avail_numel = zero2_align(avail_numel)
+
+        if debug:
+            print(f"aligned  offset={offset}, avail_numel={avail_numel}")
+
+        # Sanity check
+        if offset != avail_numel:
+            raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero2_checkpoint(world_size, fp32_flat_groups, zero_model_states):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    _zero2_merge_frozen_params(state_dict, zero_model_states)
+
+    _zero2_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def zero3_partitioned_param_info(unpartitioned_numel, world_size):
+    remainder = unpartitioned_numel % world_size
+    padding_numel = (world_size - remainder) if remainder else 0
+    partitioned_numel = math.ceil(unpartitioned_numel / world_size)
+    return partitioned_numel, padding_numel
+
+
+def _zero3_merge_frozen_params(state_dict, world_size, zero_model_states):
+    if zero_model_states[0].frozen_param_shapes is None or len(zero_model_states[0].frozen_param_shapes) == 0:
+        return
+
+    if debug:
+        for i in range(world_size):
+            num_elem = sum(s.numel() for s in zero_model_states[i].frozen_param_fragments.values())
+            print(f'rank {i}: {FROZEN_PARAM_SHAPES}.numel = {num_elem}')
+
+        frozen_param_shapes = zero_model_states[0].frozen_param_shapes
+        wanted_params = len(frozen_param_shapes)
+        wanted_numel = sum(s.numel() for s in frozen_param_shapes.values())
+        avail_numel = sum([p.numel() for p in zero_model_states[0].frozen_param_fragments.values()]) * world_size
+        print(f'Frozen params: Have {avail_numel} numels to process.')
+        print(f'Frozen params: Need {wanted_numel} numels in {wanted_params} params')
+
+    total_params = 0
+    total_numel = 0
+    for name, shape in tqdm(zero_model_states[0].frozen_param_shapes.items()):
+        total_params += 1
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+
+        param_frags = tuple(model_state.frozen_param_fragments[name] for model_state in zero_model_states)
+        state_dict[name] = torch.cat(param_frags, 0).narrow(0, 0, unpartitioned_numel).view(shape)
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Frozen params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+    print(f"Reconstructed Frozen fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states):
+    param_shapes = zero_model_states[0].param_shapes
+    avail_numel = fp32_flat_groups[0].numel() * world_size
+    # Reconstruction protocol: For zero3 we need to zip the partitions together at boundary of each
+    # param, re-consolidating each param, while dealing with padding if any
+
+    # merge list of dicts, preserving order
+    param_shapes = {k: v for d in param_shapes for k, v in d.items()}
+
+    if debug:
+        for i in range(world_size):
+            print(f"{FP32_FLAT_GROUPS}[{i}].shape={fp32_flat_groups[i].shape}")
+
+        wanted_params = len(param_shapes)
+        wanted_numel = sum(shape.numel() for shape in param_shapes.values())
+        # not asserting if there is a mismatch due to possible padding
+        avail_numel = fp32_flat_groups[0].numel() * world_size
+        print(f"Trainable params: Have {avail_numel} numels to process.")
+        print(f"Trainable params: Need {wanted_numel} numels in {wanted_params} params.")
+
+    # params
+    # XXX: for huge models that can't fit into the host's RAM we will have to recode this to support
+    # out-of-core computing solution
+    offset = 0
+    total_numel = 0
+    total_params = 0
+    for name, shape in tqdm(param_shapes.items()):
+
+        unpartitioned_numel = shape.numel()
+        total_numel += unpartitioned_numel
+        total_params += 1
+
+        partitioned_numel, partitioned_padding_numel = zero3_partitioned_param_info(unpartitioned_numel, world_size)
+
+        if debug:
+            print(
+                f"Trainable params: {total_params} {name} full shape: {shape} partition0 numel={partitioned_numel} partitioned_padding_numel={partitioned_padding_numel}"
+            )
+
+        # XXX: memory usage doubles here
+        state_dict[name] = torch.cat(
+            tuple(fp32_flat_groups[i].narrow(0, offset, partitioned_numel) for i in range(world_size)),
+            0).narrow(0, 0, unpartitioned_numel).view(shape)
+        offset += partitioned_numel
+
+    offset *= world_size
+
+    # Sanity check
+    if offset != avail_numel:
+        raise ValueError(f"consumed {offset} numels out of {avail_numel} - something is wrong")
+
+    print(f"Reconstructed Trainable fp32 state dict with {total_params} params {total_numel} elements")
+
+
+def _get_fp32_state_dict_from_zero3_checkpoint(world_size, fp32_flat_groups, zero_model_states):
+    state_dict = OrderedDict()
+
+    # buffers
+    buffers = zero_model_states[0].buffers
+    state_dict.update(buffers)
+    if debug:
+        print(f"added {len(buffers)} buffers")
+
+    _zero3_merge_frozen_params(state_dict, world_size, zero_model_states)
+
+    _zero3_merge_trainable_params(state_dict, world_size, fp32_flat_groups, zero_model_states)
+
+    # recover shared parameters
+    for pair in zero_model_states[0].shared_params:
+        state_dict[pair[0]] = state_dict[pair[1]]
+
+    return state_dict
+
+
+def get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated state_dict that can be loaded with
+    ``load_state_dict()`` and used for training without DeepSpeed or shared with others, for example
+    via a model hub.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in 'latest' file. e.g., ``global_step14``
+
+    Returns:
+        - pytorch ``state_dict``
+
+    Note: this approach may not work if your application doesn't have sufficient free CPU memory and
+    you may need to use the offline approach using the ``zero_to_fp32.py`` script that is saved with
+    the checkpoint.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import get_fp32_state_dict_from_zero_checkpoint
+        # do the training and checkpoint saving
+        state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir) # already on cpu
+        model = model.cpu() # move to cpu
+        model.load_state_dict(state_dict)
+        # submit to model hub or save the model to share with others
+
+    In this example the ``model`` will no longer be usable in the deepspeed context of the same
+    application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    If you want it all done for you, use ``load_state_dict_from_zero_checkpoint`` instead.
+
+    """
+    if tag is None:
+        latest_path = os.path.join(checkpoint_dir, 'latest')
+        if os.path.isfile(latest_path):
+            with open(latest_path, 'r') as fd:
+                tag = fd.read().strip()
+        else:
+            raise ValueError(f"Unable to find 'latest' file at {latest_path}")
+
+    ds_checkpoint_dir = os.path.join(checkpoint_dir, tag)
+
+    if not os.path.isdir(ds_checkpoint_dir):
+        raise FileNotFoundError(f"Directory '{ds_checkpoint_dir}' doesn't exist")
+
+    return _get_fp32_state_dict_from_zero_checkpoint(ds_checkpoint_dir)
+
+
+def convert_zero_checkpoint_to_fp32_state_dict(checkpoint_dir, output_file, tag=None):
+    """
+    Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict`` file that can be
+    loaded with ``torch.load(file)`` + ``load_state_dict()`` and used for training without DeepSpeed.
+
+    Args:
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``output_file``: path to the pytorch fp32 state_dict output file (e.g. path/pytorch_model.bin)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+    """
+
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+    print(f"Saving fp32 state dict to {output_file}")
+    torch.save(state_dict, output_file)
+
+
+def load_state_dict_from_zero_checkpoint(model, checkpoint_dir, tag=None):
+    """
+    1. Put the provided model to cpu
+    2. Convert ZeRO 2 or 3 checkpoint into a single fp32 consolidated ``state_dict``
+    3. Load it into the provided model
+
+    Args:
+        - ``model``: the model object to update
+        - ``checkpoint_dir``: path to the desired checkpoint folder. (one that contains the tag-folder, like ``global_step14``)
+        - ``tag``: checkpoint tag used as a unique identifier for checkpoint. If not provided will attempt to load tag in the file named ``latest`` in the checkpoint folder, e.g., ``global_step14``
+
+    Returns:
+        - ``model`: modified model
+
+    Make sure you have plenty of CPU memory available before you call this function. If you don't
+    have enough use the ``zero_to_fp32.py`` utility to do the conversion. You will find it
+    conveniently placed for you in the checkpoint folder.
+
+    A typical usage might be ::
+
+        from deepspeed.utils.zero_to_fp32 import load_state_dict_from_zero_checkpoint
+        model = load_state_dict_from_zero_checkpoint(trainer.model, checkpoint_dir)
+        # submit to model hub or save the model to share with others
+
+    Note, that once this was run, the ``model`` will no longer be usable in the deepspeed context
+    of the same application. i.e. you will need to re-initialize the deepspeed engine, since
+    ``model.load_state_dict(state_dict)`` will remove all the deepspeed magic from it.
+
+    """
+    logger.info(f"Extracting fp32 weights")
+    state_dict = get_fp32_state_dict_from_zero_checkpoint(checkpoint_dir, tag)
+
+    logger.info(f"Overwriting model with fp32 weights")
+    model = model.cpu()
+    model.load_state_dict(state_dict, strict=False)
+
+    return model
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("checkpoint_dir",
+                        type=str,
+                        help="path to the desired checkpoint folder, e.g., path/checkpoint-12")
+    parser.add_argument(
+        "output_file",
+        type=str,
+        help="path to the pytorch fp32 state_dict output file (e.g. path/checkpoint-12/pytorch_model.bin)")
+    parser.add_argument("-d", "--debug", action='store_true', help="enable debug")
+    args = parser.parse_args()
+
+    debug = args.debug
+
+    convert_zero_checkpoint_to_fp32_state_dict(args.checkpoint_dir, args.output_file)

data.py

@@ -0,0 +1,844 @@
+import copy
+import random
+import argparse
+import os
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset
+from tqdm import tqdm
+from collections import defaultdict
+import torch.distributed as dist
+import logging
+import re
+import pdb
+import json
+from prompt import sft_prompt, all_prompt
+import numpy as np
+
+
+class BaseDataset(Dataset):
+
+    def __init__(self, args):
+        super().__init__()
+
+        self.args = args
+        self.dataset = args.dataset
+        self.data_path = os.path.join(args.data_path, self.dataset)
+
+        self.max_his_len = args.max_his_len
+        self.his_sep = args.his_sep
+        self.index_file = args.index_file
+        self.add_prefix = args.add_prefix
+
+        self.new_tokens = None
+        self.allowed_tokens = None
+        self.all_items = None
+
+
+    def _load_data(self):
+
+        with open(os.path.join(self.data_path, self.dataset + self.index_file), 'r') as f:
+            self.indices = json.load(f)
+
+    def get_new_tokens(self):
+
+        if self.new_tokens is not None:
+            return self.new_tokens
+
+        self.new_tokens = set()
+        for index in self.indices.values():
+            for token in index:
+                self.new_tokens.add(token)
+        self.new_tokens = sorted(list(self.new_tokens))
+
+        return self.new_tokens
+
+    def get_all_items(self):
+
+        if self.all_items is not None:
+            return self.all_items
+
+        self.all_items = set()
+        for index in self.indices.values():
+            self.all_items.add("".join(index))
+
+        return self.all_items
+
+    def get_prefix_allowed_tokens_fn(self, tokenizer):
+
+
+        if self.allowed_tokens is None:
+            self.allowed_tokens = {}
+            for index in self.indices.values():
+                for i, token in enumerate(index):
+                    token_id = tokenizer(token)["input_ids"][1]
+                    if i not in self.allowed_tokens.keys():
+                        self.allowed_tokens[i] = set()
+                    self.allowed_tokens[i].add(token_id)
+            self.allowed_tokens[len(self.allowed_tokens.keys())] = set([tokenizer.eos_token_id])
+        sep = tokenizer("Response:")["input_ids"][1:]
+
+        def prefix_allowed_tokens_fn(batch_id, sentence):
+            sentence = sentence.tolist()
+            reversed_sent = sentence[::-1]
+            for i in range(len(reversed_sent)):
+                if reversed_sent[i:i + len(sep)] == sep[::-1]:
+                    # print(list(self.allowed_tokens[i]))
+                    return list(self.allowed_tokens[i])
+
+        return prefix_allowed_tokens_fn
+
+    def _process_data(self):
+
+        raise NotImplementedError
+
+
+
+class SeqRecDataset(BaseDataset):
+        
+    def __init__(self, args, mode="train",
+                 prompt_sample_num=1, prompt_id=0, sample_num=-1):
+        super().__init__(args)
+
+        self.mode = mode
+        self.prompt_sample_num = prompt_sample_num
+        self.prompt_id = prompt_id
+        self.sample_num = sample_num
+
+        self.prompts = all_prompt["seqrec"]
+
+
+        # load data
+        self._load_data()
+        self._remap_items()
+        
+        # load data
+        if self.mode == 'train':
+            self.inter_data = self._process_train_data()
+        elif self.mode == 'valid':
+            self.sample_valid = args.sample_valid
+            self.valid_prompt_id = args.valid_prompt_id
+            self.inter_data = self._process_valid_data()
+            self._construct_valid_text()
+        elif self.mode == 'test':
+            self.inter_data = self._process_test_data()
+        else:
+            raise NotImplementedError
+
+
+
+    def _load_data(self):
+
+        with open(os.path.join(self.data_path, self.dataset + ".inter.json"), 'r') as f:
+            self.inters = json.load(f)
+        with open(self.index_file, 'r') as f:
+            self.indices = json.load(f)
+
+
+    def _remap_items(self):
+
+        self.remapped_inters = dict()
+        for uid, items in self.inters.items():
+            new_items = ["".join(self.indices[str(i)]) for i in items]
+            self.remapped_inters[uid] = new_items
+
+
+    def _process_train_data(self):
+
+        inter_data = []
+        for uid  in self.remapped_inters:
+            items = self.remapped_inters[uid][:-2]
+            for i in range(1, len(items)):
+                one_data = dict()
+                # one_data["user"] = uid
+                one_data["item"] = items[i]
+                history = items[:i]
+                if self.max_his_len > 0:
+                    history = history[-self.max_his_len:]
+                if self.add_prefix:
+                    history = [str(k+1) + ". " + item_idx for k, item_idx in enumerate(history)]
+                one_data["inters"] = self.his_sep.join(history)
+                inter_data.append(one_data)
+
+        return inter_data
+    
+    def _process_valid_data(self):
+
+        inter_data = []
+        for uid in self.remapped_inters:
+            items = self.remapped_inters[uid]
+            one_data = dict()
+            # one_data["user"] = uid
+            one_data["item"] = items[-2]
+            history = items[:-2]
+            if self.max_his_len > 0:
+                history = history[-self.max_his_len:]
+            if self.add_prefix:
+                history = [str(k + 1) + ". " + item_idx for k, item_idx in enumerate(history)]
+            one_data["inters"] = self.his_sep.join(history)
+            inter_data.append(one_data)
+
+        return inter_data
+
+    def _process_test_data(self):
+
+        inter_data = []
+        for uid in self.remapped_inters:
+            items = self.remapped_inters[uid]
+            one_data = dict()
+            # one_data["user"] = uid
+            one_data["item"] = items[-1]
+            history = items[:-1]
+            if self.max_his_len > 0:
+                history = history[-self.max_his_len:]
+            if self.add_prefix:
+                history = [str(k + 1) + ". " + item_idx for k, item_idx in enumerate(history)]
+            one_data["inters"] = self.his_sep.join(history)
+            inter_data.append(one_data)
+
+        if self.sample_num > 0:
+            all_inter_idx = range(len(inter_data))
+            sample_idx = np.random.choice(all_inter_idx, self.sample_num, replace=False)
+            inter_data = np.array(inter_data)[sample_idx].tolist()
+
+        return inter_data
+
+    def set_prompt(self, prompt_id):
+
+        self.prompt_id = prompt_id
+
+    def __len__(self):
+        if self.mode == 'train':
+            return len(self.inter_data) * self.prompt_sample_num
+        elif self.mode == 'valid':
+            return len(self.valid_text_data)
+        elif self.mode == 'test':
+            return len(self.inter_data)
+        else:
+            raise NotImplementedError
+                    
+    def _construct_valid_text(self):
+        self.valid_text_data = []
+        if self.sample_valid:
+            all_prompt_ids = range(len(self.prompts))
+            for i in range(len(self.inter_data)):
+                d = self.inter_data[i]
+                prompt_ids = np.random.choice(all_prompt_ids, self.prompt_sample_num, replace=False)
+                for prompt_id in prompt_ids:
+                    prompt = self.prompts[prompt_id]
+                    input, output = self._get_text_data(d, prompt)
+                    self.valid_text_data.append({"input_ids": input, "labels": output})
+        else:
+            self.prompt_sample_num = 1
+            prompt = self.prompts[self.valid_prompt_id]
+            for i in range(len(self.inter_data)):
+                d = self.inter_data[i]
+                input, output = self._get_text_data(d, prompt)
+                self.valid_text_data.append({"input_ids": input, "labels": output})
+
+    def _get_text_data(self, data, prompt):
+
+        instruction = prompt["instruction"].format(**data)
+        response = prompt["response"].format(**data)
+
+        input = sft_prompt.format(instruction = instruction, response = "")
+        output = sft_prompt.format(instruction = instruction, response = response)
+
+        if self.mode == 'test':
+            return input, response
+
+        return input, output
+
+    def __getitem__(self, index):
+
+        if self.mode == 'valid':
+            return self.valid_text_data[index]
+
+        idx = index // self.prompt_sample_num
+        d = self.inter_data[idx]
+        # print(index, idx)
+
+        if self.mode == 'train':
+            prompt_id = random.randint(0, len(self.prompts) - 1)
+        elif self.mode == 'test':
+            prompt_id = self.prompt_id
+
+        prompt = self.prompts[prompt_id]
+
+        input, output = self._get_text_data(d, prompt)
+
+        # print({"input": input, "output": output})
+
+        return dict(input_ids=input, labels=output)
+
+
+class FusionSeqRecDataset(BaseDataset):
+
+    def __init__(self, args, mode="train",
+                 prompt_sample_num=1, prompt_id=0, sample_num=-1):
+        super().__init__(args)
+
+        self.mode = mode
+        self.prompt_sample_num = prompt_sample_num
+        self.prompt_id = prompt_id
+        self.sample_num = sample_num
+
+        self.prompts = all_prompt["fusionseqrec"]
+
+        # load data
+        self._load_data()
+        # self._remap_items()
+
+        # load data
+        if self.mode == 'train':
+            self.inter_data = self._process_train_data()
+        elif self.mode == 'valid':
+            self.sample_valid = args.sample_valid
+            self.valid_prompt_id = args.valid_prompt_id
+            self.inter_data = self._process_valid_data()
+            self._construct_valid_text()
+        elif self.mode == 'test':
+            self.inter_data = self._process_test_data()
+        else:
+            raise NotImplementedError
+
+
+    def _load_data(self):
+
+        with open(os.path.join(self.data_path, self.dataset + ".inter.json"), 'r') as f:
+            self.inters = json.load(f)
+        with open(self.index_file, 'r') as f:
+            self.indices = json.load(f)
+        with open(os.path.join(self.data_path, self.dataset + ".item.json"), 'r') as f:
+            self.item_feat = json.load(f)
+
+    def _process_train_data(self):
+
+        inter_data = []
+        for uid in self.inters:
+            items = self.inters[uid][:-2]
+            for i in range(1, len(items)):
+                one_data = dict()
+                # one_data["user"] = uid
+                one_data["item"] = "".join(self.indices[str(items[i])])
+                one_data["title"] = self.item_feat[str(items[i])]["title"].strip().strip(".!?,;:`")
+                one_data["description"] = self.item_feat[str(items[i])]["description"]
+                history = items[:i]
+                if self.max_his_len > 0:
+                    history = history[-self.max_his_len:]
+                inters = ["".join(self.indices[str(j)]) for j in history]
+                inter_titles = ["\"" + self.item_feat[str(j)]["title"].strip().strip(".!?,;:`") + "\"" for j in history]
+
+
+                if self.add_prefix:
+                    inters = [str(k + 1) + ". " + item_idx for k, item_idx in enumerate(inters)]
+                    inter_titles = [str(k + 1) + ". " + item_title for k, item_title in enumerate(inter_titles)]
+
+                one_data["inters"] = self.his_sep.join(inters)
+                one_data["inter_titles"] = self.his_sep.join(inter_titles)
+                inter_data.append(one_data)
+
+        if self.sample_num > 0:
+            all_inter_idx = range(len(inter_data))
+            sample_idx = np.random.choice(all_inter_idx, self.sample_num, replace=False)
+            inter_data = np.array(inter_data)[sample_idx].tolist()
+
+        return inter_data
+
+    def _process_valid_data(self):
+
+        inter_data = []
+        for uid in self.inters:
+            items = self.inters[uid]
+            one_data = dict()
+            one_data["item"] = "".join(self.indices[str(items[-2])])
+            one_data["title"] = self.item_feat[str(items[-2])]["title"].strip().strip(".!?,;:`")
+            one_data["description"] = self.item_feat[str(items[-2])]["description"]
+
+
+            history = items[:-2]
+            if self.max_his_len > 0:
+                history = history[-self.max_his_len:]
+            inters = ["".join(self.indices[str(j)]) for j in history]
+            inter_titles = ["\"" + self.item_feat[str(j)]["title"].strip().strip(".!?,;:`") + "\"" for j in history]
+
+            if self.add_prefix:
+                inters = [str(k + 1) + ". " + item_idx for k, item_idx in enumerate(inters)]
+                inter_titles = [str(k + 1) + ". " + item_title for k, item_title in enumerate(inter_titles)]
+
+            one_data["inters"] = self.his_sep.join(inters)
+            one_data["inter_titles"] = self.his_sep.join(inter_titles)
+            inter_data.append(one_data)
+
+        if self.sample_num > 0:
+            all_inter_idx = range(len(inter_data))
+            sample_idx = np.random.choice(all_inter_idx, self.sample_num, replace=False)
+            inter_data = np.array(inter_data)[sample_idx].tolist()
+
+        return inter_data
+
+    def _process_test_data(self):
+
+        inter_data = []
+        for uid in self.inters:
+            items = self.inters[uid]
+            one_data = dict()
+            one_data["item"] = "".join(self.indices[str(items[-1])])
+            one_data["title"] = self.item_feat[str(items[-1])]["title"].strip().strip(".!?,;:`")
+            one_data["description"] = self.item_feat[str(items[-1])]["description"]
+
+            history = items[:-1]
+            if self.max_his_len > 0:
+                history = history[-self.max_his_len:]
+            inters = ["".join(self.indices[str(j)]) for j in history]
+            inter_titles = ["\"" + self.item_feat[str(j)]["title"].strip().strip(".!?,;:`") + "\"" for j in history]
+
+            if self.add_prefix:
+                inters = [str(k + 1) + ". " + item_idx for k, item_idx in enumerate(inters)]
+                inter_titles = [str(k + 1) + ". " + item_title for k, item_title in enumerate(inter_titles)]
+
+            one_data["inters"] = self.his_sep.join(inters)
+            one_data["inter_titles"] = self.his_sep.join(inter_titles)
+            inter_data.append(one_data)
+
+        if self.sample_num > 0:
+            all_inter_idx = range(len(inter_data))
+            sample_idx = np.random.choice(all_inter_idx, self.sample_num, replace=False)
+            inter_data = np.array(inter_data)[sample_idx].tolist()
+
+        return inter_data
+
+    def set_prompt(self, prompt_id):
+
+        self.prompt_id = prompt_id
+
+    def __len__(self):
+        if self.mode == 'train':
+            return len(self.inter_data) * self.prompt_sample_num
+        elif self.mode == 'valid':
+            return len(self.valid_text_data)
+        elif self.mode == 'test':
+            return len(self.inter_data)
+        else:
+            raise NotImplementedError
+
+    def _construct_valid_text(self):
+        self.valid_text_data = []
+        if self.sample_valid:
+            all_prompt_ids = range(len(self.prompts))
+            for i in range(len(self.inter_data)):
+                d = self.inter_data[i]
+                prompt_ids = np.random.choice(all_prompt_ids, self.prompt_sample_num, replace=False)
+                for prompt_id in prompt_ids:
+                    prompt = self.prompts[prompt_id]
+                    input, output = self._get_text_data(d, prompt)
+                    self.valid_text_data.append({"input_ids": input, "labels": output})
+        else:
+            self.prompt_sample_num = 1
+            prompt = self.prompts[self.valid_prompt_id]
+            for i in range(len(self.inter_data)):
+                d = self.inter_data[i]
+                input, output = self._get_text_data(d, prompt)
+                self.valid_text_data.append({"input_ids": input, "labels": output})
+
+    def _get_text_data(self, data, prompt):
+
+        instruction = prompt["instruction"].format(**data)
+        response = prompt["response"].format(**data)
+
+        input = sft_prompt.format(instruction=instruction, response="")
+        output = sft_prompt.format(instruction=instruction, response=response)
+
+        if self.mode == 'test':
+            return input, response
+
+        return input, output
+
+    def __getitem__(self, index):
+
+        if self.mode == 'valid':
+            return self.valid_text_data[index]
+
+        idx = index // self.prompt_sample_num
+        d = self.inter_data[idx]
+
+        if self.mode == 'train':
+            prompt_id = random.randint(0, len(self.prompts) - 1)
+        elif self.mode == 'test':
+            prompt_id = self.prompt_id
+
+        prompt = self.prompts[prompt_id]
+
+        input, output = self._get_text_data(d, prompt)
+
+
+        return dict(input_ids=input, labels=output)
+
+
+class ItemFeatDataset(BaseDataset):
+
+    def __init__(self, args, task="item2index", prompt_sample_num=1, sample_num=-1):
+        super().__init__(args)
+
+        self.task = task.lower()
+        self.prompt_sample_num = prompt_sample_num
+        self.sample_num = sample_num
+
+        self.prompts = all_prompt[self.task]
+
+        # load data
+        self._load_data()
+        self.feat_data = self._process_data()
+
+
+
+    def _load_data(self):
+
+        with open(self.index_file, 'r') as f:
+            self.indices = json.load(f)
+        with open(os.path.join(self.data_path, self.dataset + ".item.json"), 'r') as f:
+            self.item_feat = json.load(f)
+
+
+    def _process_data(self):
+
+        feat_data = []
+        for iid in self.item_feat:
+            feat = self.item_feat[iid]
+            index = "".join(self.indices[iid])
+            feat["item"] = index
+            feat["title"] = feat["title"].strip().strip(".!?,;:`")
+            feat_data.append(feat)
+
+        if self.sample_num > 0:
+            all_idx = range(len(feat_data))
+            sample_idx = np.random.choice(all_idx, self.sample_num, replace=False)
+
+            feat_data = np.array(feat_data)[sample_idx].tolist()
+
+        return feat_data
+
+
+    def __len__(self):
+        return len(self.feat_data) * self.prompt_sample_num
+
+    def _get_text_data(self, data, prompt):
+
+        instruction = prompt["instruction"].format(**data)
+        response = prompt["response"].format(**data)
+
+        input = sft_prompt.format(instruction = instruction, response = "")
+        output = sft_prompt.format(instruction = instruction, response = response)
+
+        return input, output
+
+    def __getitem__(self, index):
+
+        idx = index // self.prompt_sample_num
+        d = self.feat_data[idx]
+
+        prompt_id = random.randint(0, len(self.prompts) - 1)
+
+        prompt = self.prompts[prompt_id]
+
+        input, output = self._get_text_data(d, prompt)
+
+        return dict(input_ids=input, labels=output)
+
+
+class ItemSearchDataset(BaseDataset):
+
+    def __init__(self, args, mode="train",
+                 prompt_sample_num=1, prompt_id=0, sample_num=-1):
+        super().__init__(args)
+
+        self.mode = mode
+        self.prompt_sample_num = prompt_sample_num
+        self.prompt_id = prompt_id
+        self.sample_num = sample_num
+
+        self.prompts = all_prompt["itemsearch"]
+
+        # load data
+        self._load_data()
+        self.search_data = self._process_data()
+
+
+
+    def _load_data(self):
+
+        with open(self.index_file, 'r') as f:
+            self.indices = json.load(f)
+        with open(os.path.join(self.data_path, self.dataset + ".user.json"), 'r') as f:
+            self.user_info = json.load(f)
+
+
+    def _process_data(self):
+
+        search_data = []
+        user_explicit_preference = self.user_info["user_explicit_preference"]
+        user_vague_intention = self.user_info["user_vague_intention"]
+        if self.mode == 'train':
+            user_vague_intention = user_vague_intention["train"]
+        elif self.mode == 'test':
+            user_vague_intention = user_vague_intention["test"]
+        else:
+            raise NotImplementedError
+
+        for uid in user_explicit_preference.keys():
+            one_data = {}
+            user_ep = user_explicit_preference[uid]
+            user_vi = user_vague_intention[uid]["querys"]
+            one_data["explicit_preferences"] = user_ep
+            one_data["user_related_intention"] = user_vi[0]
+            one_data["item_related_intention"] = user_vi[1]
+
+            iid = user_vague_intention[uid]["item"]
+            inters = user_vague_intention[uid]["inters"]
+
+            index = "".join(self.indices[str(iid)])
+            one_data["item"] = index
+
+            if self.max_his_len > 0:
+                inters = inters[-self.max_his_len:]
+            inters = ["".join(self.indices[str(i)]) for i in inters]
+            if self.add_prefix:
+                inters = [str(k + 1) + ". " + item_idx for k, item_idx in enumerate(inters)]
+
+            one_data["inters"] = self.his_sep.join(inters)
+
+            search_data.append(one_data)
+
+        if self.sample_num > 0:
+            all_idx = range(len(search_data))
+            sample_idx = np.random.choice(all_idx, self.sample_num, replace=False)
+
+            search_data = np.array(search_data)[sample_idx].tolist()
+
+        return search_data
+
+    def set_prompt(self, prompt_id):
+        self.prompt_id = prompt_id
+
+    def __len__(self):
+        if self.mode == 'train':
+            return len(self.search_data) * self.prompt_sample_num
+        elif self.mode == 'test':
+            return len(self.search_data)
+        else:
+            return len(self.search_data)
+
+
+    def _get_text_data(self, data, prompt):
+
+        instruction = prompt["instruction"].format(**data)
+        response = prompt["response"].format(**data)
+
+        input = sft_prompt.format(instruction = instruction, response = "")
+        output = sft_prompt.format(instruction = instruction, response = response)
+
+        if self.mode == 'test':
+            return input, response
+
+        return input, output
+
+    def __getitem__(self, index):
+
+        idx = index // self.prompt_sample_num
+
+        d = self.search_data[idx]
+        if self.mode == 'train':
+            prompt_id = random.randint(0, len(self.prompts) - 1)
+        elif self.mode == 'test':
+            prompt_id = self.prompt_id
+
+        prompt = self.prompts[prompt_id]
+
+        d["explicit_preference"] = copy.deepcopy(random.choice(d["explicit_preferences"]))
+        all_querys = [d["user_related_intention"], d["item_related_intention"]]
+        d["query"] = random.choice(all_querys)
+
+        input, output = self._get_text_data(d, prompt)
+
+        return dict(input_ids=input, labels=output)
+
+
+
+class PreferenceObtainDataset(BaseDataset):
+
+    def __init__(self, args, prompt_sample_num=1, sample_num=-1):
+        super().__init__(args)
+
+        self.prompt_sample_num = prompt_sample_num
+        self.sample_num = sample_num
+
+        self.prompts = all_prompt["preferenceobtain"]
+
+        # load data
+        self._load_data()
+        self._remap_items()
+
+        self.preference_data = self._process_data()
+
+
+
+    def _load_data(self):
+
+        with open(os.path.join(self.data_path, self.dataset + ".user.json"), 'r') as f:
+            self.user_info = json.load(f)
+        with open(os.path.join(self.data_path, self.dataset + ".inter.json"), 'r') as f:
+            self.inters = json.load(f)
+        with open(self.index_file, 'r') as f:
+            self.indices = json.load(f)
+
+
+    def _remap_items(self):
+
+        self.remapped_inters = dict()
+        for uid, items in self.inters.items():
+            new_items = ["".join(self.indices[str(i)]) for i in items]
+            self.remapped_inters[uid] = new_items
+
+    def _process_data(self):
+
+        preference_data = []
+        user_explicit_preference = self.user_info["user_explicit_preference"]
+
+        for uid in user_explicit_preference.keys():
+            one_data = {}
+            inters = self.remapped_inters[uid][:-3]
+            user_ep = user_explicit_preference[uid]
+
+            if self.max_his_len > 0:
+                inters = inters[-self.max_his_len:]
+            if self.add_prefix:
+                inters = [str(k + 1) + ". " + item_idx for k, item_idx in enumerate(inters)]
+
+            one_data["explicit_preferences"] = user_ep
+            one_data["inters"] = self.his_sep.join(inters)
+
+            preference_data.append(one_data)
+
+        if self.sample_num > 0:
+            all_idx = range(len(preference_data))
+            sample_idx = np.random.choice(all_idx, self.sample_num, replace=False)
+
+            preference_data = np.array(preference_data)[sample_idx].tolist()
+
+        return preference_data
+
+    def set_prompt(self, prompt_id):
+        self.prompt_id = prompt_id
+
+    def __len__(self):
+        return len(self.preference_data) * self.prompt_sample_num
+
+
+    def _get_text_data(self, data, prompt):
+
+        instruction = prompt["instruction"].format(**data)
+        response = prompt["response"].format(**data)
+
+        input = sft_prompt.format(instruction = instruction, response = "")
+        output = sft_prompt.format(instruction = instruction, response = response)
+
+        return input, output
+
+    def __getitem__(self, index):
+
+        idx = index // self.prompt_sample_num
+
+        d = self.preference_data[idx]
+        prompt_id = random.randint(0, len(self.prompts) - 1)
+
+        prompt = self.prompts[prompt_id]
+
+        d["explicit_preference"] = copy.deepcopy(random.choice(d["explicit_preferences"]))
+
+        input, output = self._get_text_data(d, prompt)
+
+        return dict(input_ids=input, labels=output)
+
+
+
+
+
+class SeqRecTestDataset(BaseDataset):
+
+    def __init__(self, args, prompt_id=0, sample_num=-1):
+        super().__init__(args)
+
+        self.prompt_id = prompt_id
+        self.sample_num = sample_num
+
+        self.prompt = all_prompt["seqrec"][self.prompt_id]
+
+        # load data
+        self._load_data()
+        self._remap_items()
+
+        self.inter_data = self._process_test_data()
+
+    def _load_data(self):
+
+        with open(os.path.join(self.data_path, self.dataset + ".inter.json"), 'r') as f:
+            self.inters = json.load(f)
+        with open(self.index_file, 'r') as f:
+            self.indices = json.load(f)
+
+
+    def _remap_items(self):
+
+        self.remapped_inters = dict()
+        for uid, items in self.inters.items():
+            new_items = ["".join(self.indices[str(i)]) for i in items]
+            self.remapped_inters[uid] = new_items
+
+    def _process_test_data(self):
+
+        inter_data = []
+        for uid in self.remapped_inters:
+            items = self.remapped_inters[uid]
+            one_data = dict()
+            # one_data["user"] = uid
+            one_data["item"] = items[-1]
+            history = items[:-1]
+            if self.max_his_len > 0:
+                history = history[-self.max_his_len:]
+            if self.add_prefix:
+                history = [str(k + 1) + ". " + item_idx for k, item_idx in enumerate(history)]
+            one_data["inters"] = self.his_sep.join(history)
+            inter_data.append(one_data)
+
+        if self.sample_num > 0:
+            all_inter_idx = range(len(inter_data))
+            sample_idx = np.random.choice(all_inter_idx, self.sample_num, replace=False)
+
+            inter_data = np.array(inter_data)[sample_idx].tolist()
+
+        return inter_data
+
+    def set_prompt(self, prompt_id):
+        self.prompt_id = prompt_id
+
+        self.prompt = all_prompt["seqrec"][self.prompt_id]
+
+    def __len__(self):
+
+        return len(self.inter_data)
+
+    def _get_text_data(self, data, prompt):
+
+        instruction = prompt["instruction"].format(**data)
+        response = prompt["response"].format(**data)
+
+        input = sft_prompt.format(instruction=instruction, response="")
+
+        return input, response
+
+    def __getitem__(self, index):
+
+        d = self.inter_data[index]
+        input, target = self._get_text_data(d, self.prompt)
+
+        return dict(input_ids=input, labels=target)

data_process/amazon18_data_process.py

@@ -0,0 +1,299 @@
+import argparse
+import collections
+import gzip
+import html
+import json
+import os
+import random
+import re
+import torch
+from tqdm import tqdm
+import numpy as np
+from utils import check_path, clean_text, amazon18_dataset2fullname, write_json_file, write_remap_index
+
+def load_ratings(file):
+    users, items, inters = set(), set(), set()
+    with open(file, 'r') as fp:
+        for line in tqdm(fp, desc='Load ratings'):
+            try:
+                item, user, rating, time = line.strip().split(',')
+                users.add(user)
+                items.add(item)
+                inters.add((user, item, float(rating), int(time)))
+            except ValueError:
+                print(line)
+    return users, items, inters
+
+
+def load_meta_items(file):
+    items = {}
+    with gzip.open(file, "r") as fp:
+        for line in tqdm(fp, desc="Load metas"):
+            data = json.loads(line)
+            item = data["asin"]
+            title = clean_text(data["title"])
+
+            descriptions = data["description"]
+            descriptions = clean_text(descriptions)
+
+            brand = data["brand"].replace("by\n", "").strip()
+
+            categories = data["category"]
+            new_categories = []
+            for category in categories:
+                if "</span>" in category:
+                    break
+                new_categories.append(category.strip())
+            categories = ",".join(new_categories).strip()
+
+            items[item] = {"title": title, "description": descriptions, "brand": brand, "categories": categories}
+            # print(items[item])
+    return items
+
+
+def load_review_data(args, user2id, item2id):
+
+    dataset_full_name = amazon18_dataset2fullname[args.dataset]
+    review_file_path = os.path.join(args.input_path, 'Review', dataset_full_name + '.json.gz')
+
+    reviews = {}
+
+    with gzip.open(review_file_path, "r") as fp:
+
+        for line in tqdm(fp,desc='Load reviews'):
+            inter = json.loads(line)
+            try:
+                user = inter['reviewerID']
+                item = inter['asin']
+                if user in user2id and item in item2id:
+                    uid = user2id[user]
+                    iid = item2id[item]
+                else:
+                    continue
+                if 'reviewText' in inter:
+                    review = clean_text(inter['reviewText'])
+                else:
+                    review = ''
+                if 'summary' in inter:
+                    summary = clean_text(inter['summary'])
+                else:
+                    summary = ''
+                reviews[str((uid,iid))]={"review":review, "summary":summary}
+
+            except ValueError:
+                print(line)
+
+    return reviews
+
+
+def get_user2count(inters):
+    user2count = collections.defaultdict(int)
+    for unit in inters:
+        user2count[unit[0]] += 1
+    return user2count
+
+
+def get_item2count(inters):
+    item2count = collections.defaultdict(int)
+    for unit in inters:
+        item2count[unit[1]] += 1
+    return item2count
+
+
+def generate_candidates(unit2count, threshold):
+    cans = set()
+    for unit, count in unit2count.items():
+        if count >= threshold:
+            cans.add(unit)
+    return cans, len(unit2count) - len(cans)
+
+
+def filter_inters(inters, can_items=None,
+                  user_k_core_threshold=0, item_k_core_threshold=0):
+    new_inters = []
+
+    # filter by meta items
+    if can_items:
+        print('\nFiltering by meta items: ')
+        for unit in inters:
+            if unit[1] in can_items.keys():
+                new_inters.append(unit)
+        inters, new_inters = new_inters, []
+        print('    The number of inters: ', len(inters))
+
+    # filter by k-core
+    if user_k_core_threshold or item_k_core_threshold:
+        print('\nFiltering by k-core:')
+        idx = 0
+        user2count = get_user2count(inters)
+        item2count = get_item2count(inters)
+
+        while True:
+            new_user2count = collections.defaultdict(int)
+            new_item2count = collections.defaultdict(int)
+            users, n_filtered_users = generate_candidates( # users is set
+                user2count, user_k_core_threshold)
+            items, n_filtered_items = generate_candidates(
+                item2count, item_k_core_threshold)
+            if n_filtered_users == 0 and n_filtered_items == 0:
+                break
+            for unit in inters:
+                if unit[0] in users and unit[1] in items:
+                    new_inters.append(unit)
+                    new_user2count[unit[0]] += 1
+                    new_item2count[unit[1]] += 1
+            idx += 1
+            inters, new_inters = new_inters, []
+            user2count, item2count = new_user2count, new_item2count
+            print('    Epoch %d The number of inters: %d, users: %d, items: %d'
+                    % (idx, len(inters), len(user2count), len(item2count)))
+    return inters
+
+
+def make_inters_in_order(inters):
+    user2inters, new_inters = collections.defaultdict(list), list()
+    for inter in inters:
+        user, item, rating, timestamp = inter
+        user2inters[user].append((user, item, rating, timestamp))
+    for user in user2inters:
+        user_inters = user2inters[user]
+        user_inters.sort(key=lambda d: d[3])
+        interacted_item = set()
+        for inter in user_inters:
+            if inter[1] in interacted_item: # 过滤重复交互
+                continue
+            interacted_item.add(inter[1])
+            new_inters.append(inter)
+    return new_inters
+
+
+def preprocess_rating(args):
+    dataset_full_name = amazon18_dataset2fullname[args.dataset]
+
+    print('Process rating data: ')
+    print(' Dataset: ', args.dataset)
+
+    # load ratings
+    rating_file_path = os.path.join(args.input_path, 'Ratings', dataset_full_name + '.csv')
+    rating_users, rating_items, rating_inters = load_ratings(rating_file_path)
+
+    # load item IDs with meta data
+    meta_file_path = os.path.join(args.input_path, 'Metadata', f'meta_{dataset_full_name}.json.gz')
+    meta_items = load_meta_items(meta_file_path)
+
+    # 1. Filter items w/o meta data;
+    # 2. K-core filtering;
+    print('The number of raw inters: ', len(rating_inters))
+
+    rating_inters = make_inters_in_order(rating_inters)
+
+    rating_inters = filter_inters(rating_inters, can_items=meta_items,
+                                  user_k_core_threshold=args.user_k,
+                                  item_k_core_threshold=args.item_k)
+
+    # sort interactions chronologically for each user
+    rating_inters = make_inters_in_order(rating_inters)
+    print('\n')
+
+    # return: list of (user_ID, item_ID, rating, timestamp)
+    return rating_inters, meta_items
+
+def convert_inters2dict(inters):
+    user2items = collections.defaultdict(list)
+    user2index, item2index = dict(), dict()
+    for inter in inters:
+        user, item, rating, timestamp = inter
+        if user not in user2index:
+            user2index[user] = len(user2index)
+        if item not in item2index:
+            item2index[item] = len(item2index)
+        user2items[user2index[user]].append(item2index[item])
+    return user2items, user2index, item2index
+
+def generate_data(args, rating_inters):
+    print('Split dataset: ')
+    print(' Dataset: ', args.dataset)
+
+    # generate train valid temp
+    user2items, user2index, item2index = convert_inters2dict(rating_inters)
+    train_inters, valid_inters, test_inters = dict(), dict(), dict()
+    for u_index in range(len(user2index)):
+        inters = user2items[u_index]
+        # leave one out
+        train_inters[u_index] = [str(i_index) for i_index in inters[:-2]]
+        valid_inters[u_index] = [str(inters[-2])]
+        test_inters[u_index] = [str(inters[-1])]
+        assert len(user2items[u_index]) == len(train_inters[u_index]) + \
+               len(valid_inters[u_index]) + len(test_inters[u_index])
+    return user2items, train_inters, valid_inters, test_inters, user2index, item2index
+
+def convert_to_atomic_files(args, train_data, valid_data, test_data):
+    print('Convert dataset: ')
+    print(' Dataset: ', args.dataset)
+    uid_list = list(train_data.keys())
+    uid_list.sort(key=lambda t: int(t))
+
+    with open(os.path.join(args.output_path, args.dataset, f'{args.dataset}.train.inter'), 'w') as file:
+        file.write('user_id:token\titem_id_list:token_seq\titem_id:token\n')
+        for uid in uid_list:
+            item_seq = train_data[uid]
+            seq_len = len(item_seq)
+            for target_idx in range(1, seq_len):
+                target_item = item_seq[-target_idx]
+                seq = item_seq[:-target_idx][-50:]
+                file.write(f'{uid}\t{" ".join(seq)}\t{target_item}\n')
+
+    with open(os.path.join(args.output_path, args.dataset, f'{args.dataset}.valid.inter'), 'w') as file:
+        file.write('user_id:token\titem_id_list:token_seq\titem_id:token\n')
+        for uid in uid_list:
+            item_seq = train_data[uid][-50:]
+            target_item = valid_data[uid][0]
+            file.write(f'{uid}\t{" ".join(item_seq)}\t{target_item}\n')
+
+    with open(os.path.join(args.output_path, args.dataset, f'{args.dataset}.test.inter'), 'w') as file:
+        file.write('user_id:token\titem_id_list:token_seq\titem_id:token\n')
+        for uid in uid_list:
+            item_seq = (train_data[uid] + valid_data[uid])[-50:]
+            target_item = test_data[uid][0]
+            file.write(f'{uid}\t{" ".join(item_seq)}\t{target_item}\n')
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--dataset', type=str, default='Arts', help='Instruments / Arts / Games')
+    parser.add_argument('--user_k', type=int, default=5, help='user k-core filtering')
+    parser.add_argument('--item_k', type=int, default=5, help='item k-core filtering')
+    parser.add_argument('--input_path', type=str, default='')
+    parser.add_argument('--output_path', type=str, default='')
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    args = parse_args()
+
+    # load interactions from raw rating file
+    rating_inters, meta_items = preprocess_rating(args)
+
+
+    # split train/valid/temp
+    all_inters,train_inters, valid_inters, test_inters, user2index, item2index = generate_data(args, rating_inters)
+
+    check_path(os.path.join(args.output_path, args.dataset))
+
+    write_json_file(all_inters, os.path.join(args.output_path, args.dataset, f'{args.dataset}.inter.json'))
+    convert_to_atomic_files(args, train_inters, valid_inters, test_inters)
+
+    item2feature = collections.defaultdict(dict)
+    for item, item_id in item2index.items():
+        item2feature[item_id] = meta_items[item]
+
+    # reviews = load_review_data(args, user2index, item2index)
+
+    print("user:",len(user2index))
+    print("item:",len(item2index))
+
+    write_json_file(item2feature, os.path.join(args.output_path, args.dataset, f'{args.dataset}.item.json'))
+    # write_json_file(reviews, os.path.join(args.output_path, args.dataset, f'{args.dataset}.review.json'))
+
+
+    write_remap_index(user2index, os.path.join(args.output_path, args.dataset, f'{args.dataset}.user2id'))
+    write_remap_index(item2index, os.path.join(args.output_path, args.dataset, f'{args.dataset}.item2id'))

data_process/amazon18_recbole_data_process.py

@@ -0,0 +1,226 @@
+import argparse
+import collections
+import gzip
+import html
+import json
+import os
+import random
+import re
+import torch
+from tqdm import tqdm
+import numpy as np
+from utils import check_path, clean_text, amazon18_dataset2fullname,write_json_file,write_remap_index
+
+def load_ratings(file):
+    users, items, inters = set(), set(), set()
+    with open(file, 'r') as fp:
+        for line in tqdm(fp, desc='Load ratings'):
+            try:
+                item, user, rating, time = line.strip().split(',')
+                users.add(user)
+                items.add(item)
+                inters.add((user, item, float(rating), int(time)))
+            except ValueError:
+                print(line)
+    return users, items, inters
+
+
+def load_meta_items(file):
+    items = {}
+    # re_tag = re.compile('</?\w+[^>]*>')
+    with gzip.open(file, "r") as fp:
+        for line in tqdm(fp, desc="Load metas"):
+            data = json.loads(line)
+            item = data["asin"]
+            title = clean_text(data["title"])
+
+            descriptions = data["description"]
+            descriptions = clean_text(descriptions)
+            # new_descriptions = []
+            # for description in descriptions:
+            #     description = re.sub(re_tag, '', description)
+            #     new_descriptions.append(description.strip())
+            # descriptions = " ".join(new_descriptions).strip()
+
+            brand = data["brand"].replace("by\n", "").strip()
+
+            categories = data["category"]
+            new_categories = []
+            for category in categories:
+                if "</span>" in category:
+                    break
+                new_categories.append(category.strip())
+            categories = ",".join(new_categories[1:]).strip()
+
+            items[item] = {"title": title, "description": descriptions, "brand": brand, "categories": categories}
+            # print(items[item])
+    return items
+
+
+def get_user2count(inters):
+    user2count = collections.defaultdict(int)
+    for unit in inters:
+        user2count[unit[0]] += 1
+    return user2count
+
+
+def get_item2count(inters):
+    item2count = collections.defaultdict(int)
+    for unit in inters:
+        item2count[unit[1]] += 1
+    return item2count
+
+
+def generate_candidates(unit2count, threshold):
+    cans = set()
+    for unit, count in unit2count.items():
+        if count >= threshold:
+            cans.add(unit)
+    return cans, len(unit2count) - len(cans)
+
+
+def filter_inters(inters, can_items=None,
+                  user_k_core_threshold=0, item_k_core_threshold=0):
+    new_inters = []
+
+    # filter by meta items
+    if can_items:
+        print('\nFiltering by meta items: ')
+        for unit in inters:
+            if unit[1] in can_items.keys():
+                new_inters.append(unit)
+        inters, new_inters = new_inters, []
+        print('    The number of inters: ', len(inters))
+
+    # filter by k-core
+    if user_k_core_threshold or item_k_core_threshold:
+        print('\nFiltering by k-core:')
+        idx = 0
+        user2count = get_user2count(inters)
+        item2count = get_item2count(inters)
+
+        while True:
+            new_user2count = collections.defaultdict(int)
+            new_item2count = collections.defaultdict(int)
+            users, n_filtered_users = generate_candidates( # users is set
+                user2count, user_k_core_threshold)
+            items, n_filtered_items = generate_candidates(
+                item2count, item_k_core_threshold)
+            if n_filtered_users == 0 and n_filtered_items == 0:
+                break
+            for unit in inters:
+                if unit[0] in users and unit[1] in items:
+                    new_inters.append(unit)
+                    new_user2count[unit[0]] += 1
+                    new_item2count[unit[1]] += 1
+            idx += 1
+            inters, new_inters = new_inters, []
+            user2count, item2count = new_user2count, new_item2count
+            print('    Epoch %d The number of inters: %d, users: %d, items: %d'
+                    % (idx, len(inters), len(user2count), len(item2count)))
+    return inters
+
+
+def make_inters_in_order(inters):
+    user2inters, new_inters = collections.defaultdict(list), list()
+    for inter in inters:
+        user, item, rating, timestamp = inter
+        user2inters[user].append((user, item, rating, timestamp))
+    for user in user2inters:
+        user_inters = user2inters[user]
+        user_inters.sort(key=lambda d: d[3])
+        interacted_item = set()
+        for inter in user_inters:
+            if inter[1] in interacted_item: # 过滤重复交互
+                continue
+            interacted_item.add(inter[1])
+            new_inters.append(inter)
+    return new_inters
+
+
+def preprocess_rating(args):
+    dataset_full_name = amazon18_dataset2fullname[args.dataset]
+
+    print('Process rating data: ')
+    print(' Dataset: ', args.dataset)
+
+    # load ratings
+    rating_file_path = os.path.join(args.input_path, 'Ratings', dataset_full_name + '.csv')
+    rating_users, rating_items, rating_inters = load_ratings(rating_file_path)
+
+    # load item IDs with meta data
+    meta_file_path = os.path.join(args.input_path, 'Metadata', f'meta_{dataset_full_name}.json.gz')
+    meta_items = load_meta_items(meta_file_path)
+
+    # 1. Filter items w/o meta data;
+    # 2. K-core filtering;
+    print('The number of raw inters: ', len(rating_inters))
+
+    rating_inters = make_inters_in_order(rating_inters)
+
+    rating_inters = filter_inters(rating_inters, can_items=meta_items,
+                                  user_k_core_threshold=args.user_k,
+                                  item_k_core_threshold=args.item_k)
+
+    # sort interactions chronologically for each user
+    rating_inters = make_inters_in_order(rating_inters)
+    print('\n')
+
+    # return: list of (user_ID, item_ID, rating, timestamp)
+    return rating_inters, meta_items
+
+def save_inter(args, inters):
+    print('Convert dataset: ')
+    print(' Dataset: ', args.dataset)
+
+    with open(os.path.join(args.output_path, args.dataset, f'{args.dataset}.inter'), 'w') as file:
+        file.write('user_id:token\titem_id:token\trating:float\ttimestamp:float\n')
+        for inter in inters:
+            user, item, rating, timestamp = inter
+            file.write(f'{user}\t{item}\t{rating}\t{timestamp}\n')
+
+
+def save_feat(args, feat, all_items):
+    iid_list = list(feat.keys())
+    num_item = 0
+    with open(os.path.join(args.output_path, args.dataset, f'{args.dataset}.item'), 'w') as file:
+        # "title": title, "description": descriptions, "brand": brand, "categories": categories
+        file.write('item_id:token\ttitle:token_seq\tbrand:token\tcategories:token_seq\n')
+        for iid in iid_list:
+            if iid in all_items:
+                num_item += 1
+                title, brand, categories = feat[iid]["title"], feat[iid]["brand"], feat[iid]["categories"]
+                file.write(f'{iid}\t{title}\t{brand}\t{categories}\n')
+    print("num_item: ", num_item)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--dataset', type=str, default='Arts', help='Instruments / Arts / Games')
+    parser.add_argument('--user_k', type=int, default=5, help='user k-core filtering')
+    parser.add_argument('--item_k', type=int, default=5, help='item k-core filtering')
+    parser.add_argument('--input_path', type=str, default='')
+    parser.add_argument('--output_path', type=str, default='')
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    args = parse_args()
+
+    # load interactions from raw rating file
+    rating_inters, meta_items = preprocess_rating(args)
+
+    check_path(os.path.join(args.output_path, args.dataset))
+
+
+    all_items = set()
+    for inter in rating_inters:
+        user, item, rating, timestamp = inter
+        all_items.add(item)
+
+    print("total item: ", len(list(all_items)))
+
+    save_inter(args,rating_inters)
+    save_feat(args,meta_items, all_items)
+
+

data_process/amazon_text_emb.py

@@ -0,0 +1,161 @@
+import argparse
+import collections
+import gzip
+import html
+import json
+import os
+import random
+import re
+import torch
+from tqdm import tqdm
+import numpy as np
+from utils import *
+from transformers import LlamaForCausalLM, LlamaTokenizer, LlamaConfig, AutoTokenizer, AutoModel
+
+
+def load_data(args):
+
+    item2feature_path = args.data_path
+    item2feature = load_json(item2feature_path)
+
+    return item2feature
+
+def generate_text(item2feature, features):
+    item_text_list = []
+
+    for item in item2feature:
+        data = item2feature[item]
+        text = []
+        for meta_key in features:
+            if meta_key in data:
+                meta_value = clean_text(data[meta_key])
+                text.append(meta_value.strip())
+
+        item_text_list.append([int(item), text])
+
+    return item_text_list
+
+def preprocess_text(args):
+    print('Process text data ...')
+    # print('Dataset: ', args.dataset)
+
+    item2feature = load_data(args)
+    # load item text and clean
+    item_text_list = generate_text(item2feature, ['title'])
+    # item_text_list = generate_text(item2feature, ['title'])
+    # return: list of (item_ID, cleaned_item_text)
+    return item_text_list
+
+def generate_item_embedding(args, item_text_list, tokenizer, model, word_drop_ratio=-1, save_path = ''):
+    print('Generate text embedding ...')
+    # print(' Dataset: ', args.dataset)
+
+    items, texts = zip(*item_text_list)
+    order_texts = [[0]] * len(items)
+    for item, text in zip(items, texts):
+        order_texts[item] = text
+    for text in order_texts:
+        assert text != [0]
+
+    embeddings = []
+    emb_result = []
+    start, batch_size = 0, 1
+    with torch.no_grad():
+        while start < len(order_texts):
+            if (start+1)%100==0:
+                print("==>",start+1)
+            field_texts = order_texts[start: start + batch_size]
+            # print(field_texts)
+            field_texts = zip(*field_texts)
+    
+            field_embeddings = []
+            for sentences in field_texts:
+                sentences = list(sentences)
+                # print(sentences)
+                if word_drop_ratio > 0:
+                    print(f'Word drop with p={word_drop_ratio}')
+                    new_sentences = []
+                    for sent in sentences:
+                        new_sent = []
+                        sent = sent.split(' ')
+                        for wd in sent:
+                            rd = random.random()
+                            if rd > word_drop_ratio:
+                                new_sent.append(wd)
+                        new_sent = ' '.join(new_sent)
+                        new_sentences.append(new_sent)
+                    sentences = new_sentences
+                encoded_sentences = tokenizer(sentences, max_length=args.max_sent_len,
+                                              truncation=True, return_tensors='pt',padding="longest").to(args.device)
+                outputs = model(input_ids=encoded_sentences.input_ids,
+                                attention_mask=encoded_sentences.attention_mask)
+    
+                masked_output = outputs.last_hidden_state * encoded_sentences['attention_mask'].unsqueeze(-1)
+                mean_output = masked_output.sum(dim=1) / encoded_sentences['attention_mask'].sum(dim=-1, keepdim=True)
+                mean_output = mean_output.detach().cpu()
+                emb_result.append(mean_output.numpy().tolist())
+                field_embeddings.append(mean_output)
+    
+            field_mean_embedding = torch.stack(field_embeddings, dim=0).mean(dim=0)
+            embeddings.append(field_mean_embedding)
+            start += batch_size
+
+    embeddings = torch.cat(embeddings, dim=0).numpy()
+    print('Embeddings shape: ', embeddings.shape)
+
+    all_results = {
+        'text':[],
+        'node_type':[],
+        'emb':[]
+    }
+
+    all_results['text'] = [t[0] for t in texts]
+    all_results['node_type'] = [1] * len(all_results['text'])
+    for emb in emb_result:
+        str_emb = ''
+        for x in emb:
+            str_emb = str_emb + str(x) + ' '
+    all_results['emb'].append(str_emb[:-1])
+    
+    import pandas as pd
+    df = pd.DataFrame(all_results)
+    # header = 0: w/o column name; index = False: w/o index column 
+    df.to_csv(args.save_path, sep = '\t', header = 0, index = False)
+
+    # file = os.path.join(args.root, args.dataset + '.emb-' + args.plm_name + "-td" + ".npy")
+    # np.save(file, embeddings)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--dataset', type=str, default='Arts', help='Instruments / Arts / Games')
+    parser.add_argument('--root', type=str, default="")
+    parser.add_argument('--gpu_id', type=int, default=0, help='ID of running GPU')
+    parser.add_argument('--plm_name', type=str, default='llama')
+    parser.add_argument('--plm_checkpoint', type=str,
+                        default='')
+    parser.add_argument('--max_sent_len', type=int, default=2048)
+    parser.add_argument('--word_drop_ratio', type=float, default=-1, help='word drop ratio, do not drop by default')
+    parser.add_argument('--data_path', type=str, default='')
+    parser.add_argument('--save_path', type=str, default='')
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    args = parse_args()
+
+    args.root = os.path.join(args.root, args.dataset)
+
+    device = set_device(args.gpu_id)
+    args.device = device
+
+    item_text_list = preprocess_text(args)
+
+    plm_tokenizer, plm_model = load_plm(args.plm_checkpoint)
+    if plm_tokenizer.pad_token_id is None:
+        plm_tokenizer.pad_token_id = 0
+    plm_model = plm_model.to(device)
+
+    generate_item_embedding(args, item_text_list, plm_tokenizer,
+                            plm_model, word_drop_ratio = args.word_drop_ratio, 
+                            save_path = args.save_path)

data_process/get_llm_output.py

@@ -0,0 +1,374 @@
+
+
+import argparse
+import os
+import os.path as osp
+import random
+import time
+from logging import getLogger
+import openai
+from utils import get_res_batch, load_json, intention_prompt, preference_prompt_1, preference_prompt_2, amazon18_dataset2fullname, write_json_file
+import json
+
+
+
+def get_intention_train(args, inters, item2feature, reviews, api_info):
+
+    intention_train_output_file = os.path.join(args.root,"intention_train.json")
+
+
+    # Suggest modifying the prompt based on different datasets
+    prompt = intention_prompt
+    dataset_full_name = amazon18_dataset2fullname[args.dataset]
+    dataset_full_name = dataset_full_name.replace("_", " ").lower()
+    print(dataset_full_name)
+
+    prompt_list = []
+
+    inter_data = []
+
+    for (user,item_list) in inters.items():
+        user = int(user)
+        item = int(item_list[-3])
+        history = item_list[:-3]
+
+        inter_data.append((user,item,history))
+
+        review = reviews[str((user, item))]["review"]
+        item_title = item2feature[str(item)]["title"]
+        input_prompt = prompt.format(item_title=item_title,dataset_full_name=dataset_full_name,review=review)
+        prompt_list.append(input_prompt)
+
+    st = 0
+    with open(intention_train_output_file, mode='a') as f:
+
+        while st < len(prompt_list):
+        # while st < 3:
+            print(st)
+            # if st < 25631:
+            #     st += args.batchsize
+            #     continue
+
+
+            res = get_res_batch(args.model_name, prompt_list[st:st+args.batchsize], args.max_tokens, api_info)
+
+            for i, answer in enumerate(res):
+                user, item, history = inter_data[st+i]
+                # print(answer)
+                # print("=============")
+
+                if answer == '':
+                    print("answer null error")
+                    answer = "I enjoy high-quality item."
+
+                if answer.strip().count('\n') != 1:
+                    if 'haracteristics:' in answer:
+                        answer = answer.strip().split("The item's characteristics:")
+                    else:
+                        answer = answer.strip().split("The item's characteristic:")
+                else:
+                    answer = answer.strip().split('\n')
+
+                if '' in answer:
+                    answer.remove('')
+
+                if len(answer) == 1:
+                    print(answer)
+                    user_preference = item_character = answer[0]
+                elif len(answer) >= 3:
+                    print(answer)
+                    answer = answer[-1]
+                    user_preference = item_character = answer
+                else:
+                    user_preference, item_character = answer
+
+                if ':' in user_preference:
+                    idx = user_preference.index(':')
+                    user_preference = user_preference[idx+1:]
+                user_preference = user_preference.strip().replace('}','')
+                user_preference = user_preference.replace('\n','')
+
+                if ':' in item_character:
+                    idx = item_character.index(':')
+                    item_character = item_character[idx+1:]
+                item_character = item_character.strip().replace('}','')
+                item_character = item_character.replace('\n','')
+
+
+                dict = {"user":user, "item":item, "inters": history,
+                        "user_related_intention":user_preference, "item_related_intention": item_character}
+
+                json.dump(dict, f)
+                f.write("\n")
+
+            st += args.batchsize
+
+    return intention_train_output_file
+
+
+def get_intention_test(args, inters, item2feature, reviews, api_info):
+
+    intention_test_output_file = os.path.join(args.root,"intention_test.json")
+
+    # Suggest modifying the prompt based on different datasets
+    prompt = intention_prompt
+    dataset_full_name = amazon18_dataset2fullname[args.dataset]
+    dataset_full_name = dataset_full_name.replace("_", " ").lower()
+    print(dataset_full_name)
+
+    prompt_list = []
+
+    inter_data = []
+
+    for (user,item_list) in inters.items():
+        user = int(user)
+        item = int(item_list[-1])
+        history = item_list[:-1]
+
+        inter_data.append((user,item,history))
+
+        review = reviews[str((user, item))]["review"]
+        item_title = item2feature[str(item)]["title"]
+        input_prompt = prompt.format(item_title=item_title,dataset_full_name=dataset_full_name,review=review)
+        prompt_list.append(input_prompt)
+
+    st = 0
+    with open(intention_test_output_file, mode='a') as f:
+
+        while st < len(prompt_list):
+        # while st < 3:
+            print(st)
+            # if st < 4623:
+            #     st += args.batchsize
+            #     continue
+
+            res = get_res_batch(args.model_name, prompt_list[st:st+args.batchsize], args.max_tokens, api_info)
+
+            for i, answer in enumerate(res):
+                user, item, history = inter_data[st+i]
+
+                if answer == '':
+                    print("answer null error")
+                    answer = "I enjoy high-quality item."
+
+                if answer.strip().count('\n') != 1:
+                    if 'haracteristics:' in answer:
+                        answer = answer.strip().split("The item's characteristics:")
+                    else:
+                        answer = answer.strip().split("The item's characteristic:")
+                else:
+                    answer = answer.strip().split('\n')
+
+                if '' in answer:
+                    answer.remove('')
+
+                if len(answer) == 1:
+                    print(answer)
+                    user_preference = item_character = answer[0]
+                elif len(answer) >= 3:
+                    print(answer)
+                    answer = answer[-1]
+                    user_preference = item_character = answer
+                else:
+                    user_preference, item_character = answer
+
+                if ':' in user_preference:
+                    idx = user_preference.index(':')
+                    user_preference = user_preference[idx+1:]
+                user_preference = user_preference.strip().replace('}','')
+                user_preference = user_preference.replace('\n','')
+
+                if ':' in item_character:
+                    idx = item_character.index(':')
+                    item_character = item_character[idx+1:]
+                item_character = item_character.strip().replace('}','')
+                item_character = item_character.replace('\n','')
+
+
+                dict = {"user":user, "item":item, "inters": history,
+                        "user_related_intention":user_preference, "item_related_intention": item_character}
+
+                json.dump(dict, f)
+                f.write("\n")
+
+            st += args.batchsize
+
+    return intention_test_output_file
+
+
+
+
+def get_user_preference(args, inters, item2feature, reviews, api_info):
+
+    preference_output_file = os.path.join(args.root,"user_preference.json")
+
+
+    # Suggest modifying the prompt based on different datasets
+    prompt_1 = preference_prompt_1
+    prompt_2 = preference_prompt_2
+
+
+    dataset_full_name = amazon18_dataset2fullname[args.dataset]
+    dataset_full_name = dataset_full_name.replace("_", " ").lower()
+    print(dataset_full_name)
+
+    prompt_list_1 = []
+    prompt_list_2 = []
+
+    users = []
+
+    for (user,item_list) in inters.items():
+        users.append(user)
+        history = item_list[:-3]
+        item_titles = []
+        for j, item in enumerate(history):
+            item_titles.append(str(j+1) + '.' + item2feature[str(item)]["title"])
+        if len(item_titles) > args.max_his_len:
+            item_titles = item_titles[-args.max_his_len:]
+        item_titles = ", ".join(item_titles)
+        
+        input_prompt_1 = prompt_1.format(dataset_full_name=dataset_full_name, item_titles=item_titles)
+        input_prompt_2 = prompt_2.format(dataset_full_name=dataset_full_name, item_titles=item_titles)
+
+        prompt_list_1.append(input_prompt_1)
+        prompt_list_2.append(input_prompt_2)
+
+
+    st = 0
+    with open(preference_output_file, mode='a') as f:
+
+        while st < len(prompt_list_1):
+        # while st < 3:
+            print(st)
+            # if st < 22895:
+            #     st += args.batchsize
+            #     continue
+
+            res_1 = get_res_batch(args.model_name, prompt_list_1[st:st + args.batchsize], args.max_tokens, api_info)
+            res_2 = get_res_batch(args.model_name, prompt_list_2[st:st + args.batchsize], args.max_tokens, api_info)
+            for i, answers in enumerate(zip(res_1, res_2)):
+
+                user = users[st + i]
+
+                answer_1, answer_2 = answers
+                # print(answers)
+                # print("=============")
+
+                if answer_1 == '':
+                    print("answer null error")
+                    answer_1 = "I enjoy high-quality item."
+                    
+                if answer_2 == '':
+                    print("answer null error")
+                    answer_2 = "I enjoy high-quality item."
+
+                if answer_2.strip().count('\n') != 1:
+                    if 'references:' in answer_2:
+                        answer_2 = answer_2.strip().split("Short-term preferences:")
+                    else:
+                        answer_2 = answer_2.strip().split("Short-term preference:")
+                else:
+                    answer_2 = answer_2.strip().split('\n')
+
+                if '' in answer_2:
+                    answer_2.remove('')
+
+                if len(answer_2) == 1:
+                    print(answer_2)
+                    long_preference = short_preference = answer_2[0]
+                elif len(answer_2) >= 3:
+                    print(answer_2)
+                    answer_2 = answer_2[-1]
+                    long_preference = short_preference = answer_2
+                else:
+                    long_preference, short_preference = answer_2
+
+                if ':' in long_preference:
+                    idx = long_preference.index(':')
+                    long_preference = long_preference[idx+1:]
+                long_preference = long_preference.strip().replace('}','')
+                long_preference = long_preference.replace('\n','')
+
+                if ':' in short_preference:
+                    idx = short_preference.index(':')
+                    short_preference = short_preference[idx+1:]
+                short_preference = short_preference.strip().replace('}','')
+                short_preference = short_preference.replace('\n','')
+
+                dict = {"user":user,"user_preference":[answer_1, long_preference, short_preference]}
+                # print(dict)
+                json.dump(dict, f)
+                f.write("\n")
+
+            st += args.batchsize
+
+    return preference_output_file
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--dataset', type=str, default='Instruments', help='Instruments / Arts / Games')
+    parser.add_argument('--root', type=str, default='')
+    parser.add_argument('--api_info', type=str, default='./api_info.json')
+    parser.add_argument('--model_name', type=str, default='text-davinci-003')
+    parser.add_argument('--max_tokens', type=int, default=512)
+    parser.add_argument('--batchsize', type=int, default=16)
+    parser.add_argument('--max_his_len', type=int, default=20)
+    return parser.parse_args()
+
+if __name__ == "__main__":
+    args = parse_args()
+
+    args.root = os.path.join(args.root, args.dataset)
+
+    api_info = load_json(args.api_info)
+    openai.api_key = api_info["api_key_list"].pop()
+
+
+    inter_path = os.path.join(args.root, f'{args.dataset}.inter.json')
+    inters = load_json(inter_path)
+
+
+    item2feature_path = os.path.join(args.root, f'{args.dataset}.item.json')
+    item2feature = load_json(item2feature_path)
+
+    reviews_path = os.path.join(args.root, f'{args.dataset}.review.json')
+    reviews = load_json(reviews_path)
+
+    intention_train_output_file = get_intention_train(args, inters, item2feature, reviews, api_info)
+    intention_test_output_file = get_intention_test(args, inters, item2feature, reviews ,api_info)
+    preference_output_file = get_user_preference(args, inters, item2feature, reviews, api_info)
+
+    intention_train = {}
+    intention_test = {}
+    user_preference = {}
+
+    with open(intention_train_output_file, "r") as f:
+        for line in f:
+            # print(line)
+            content = json.loads(line)
+            if content["user"] not in intention_train:
+                intention_train[content["user"]] = {"item":content["item"],
+                                                "inters":content["inters"],
+                                                "querys":[ content["user_related_intention"], content["item_related_intention"] ]}
+
+
+    with open(intention_test_output_file, "r") as f:
+        for line in f:
+            content = json.loads(line)
+            if content["user"] not in intention_train:
+                intention_test[content["user"]] = {"item":content["item"],
+                                                "inters":content["inters"],
+                                                "querys":[ content["user_related_intention"], content["item_related_intention"] ]}
+
+
+    with open(preference_output_file, "r") as f:
+        for line in f:
+            content = json.loads(line)
+            user_preference[content["user"]] = content["user_preference"]
+
+    user_dict = {
+        "user_explicit_preference": user_preference,
+        "user_vague_intention": {"train": intention_train, "test": intention_test},
+    }
+
+    write_json_file(user_dict, os.path.join(args.root, f'{args.dataset}.user.json'))

data_process/utils.py

@@ -0,0 +1,238 @@
+import html
+import json
+import os
+import pickle
+import re
+import time
+
+import torch
+# import gensim
+from transformers import AutoModel, AutoTokenizer
+import collections
+import openai
+
+
+
+def get_res_batch(model_name, prompt_list, max_tokens, api_info):
+
+    while True:
+        try:
+            res = openai.Completion.create(
+                model=model_name,
+                prompt=prompt_list,
+                temperature=0.4,
+                max_tokens=max_tokens,
+                top_p=1,
+                frequency_penalty=0,
+                presence_penalty=0
+            )
+            output_list = []
+            for choice in res['choices']:
+                output = choice['text'].strip()
+                output_list.append(output)
+
+            return output_list
+
+        except openai.error.AuthenticationError as e:
+            print(e)
+            openai.api_key = api_info["api_key_list"].pop()
+            time.sleep(10)
+        except openai.error.RateLimitError as e:
+            print(e)
+            if str(e) == "You exceeded your current quota, please check your plan and billing details.":
+                openai.api_key = api_info["api_key_list"].pop()
+                time.sleep(10)
+            else:
+                print('\nopenai.error.RateLimitError\nRetrying...')
+                time.sleep(10)
+        except openai.error.ServiceUnavailableError as e:
+            print(e)
+            print('\nopenai.error.ServiceUnavailableError\nRetrying...')
+            time.sleep(10)
+        except openai.error.Timeout:
+            print('\nopenai.error.Timeout\nRetrying...')
+            time.sleep(10)
+        except openai.error.APIError as e:
+            print(e)
+            print('\nopenai.error.APIError\nRetrying...')
+            time.sleep(10)
+        except openai.error.APIConnectionError as e:
+            print(e)
+            print('\nopenai.error.APIConnectionError\nRetrying...')
+            time.sleep(10)
+        except Exception as e:
+            print(e)
+            return None
+
+
+
+
+def check_path(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+
+def set_device(gpu_id):
+    if gpu_id == -1:
+        return torch.device('cpu')
+    else:
+        return torch.device(
+            'cuda:' + str(gpu_id) if torch.cuda.is_available() else 'cpu')
+
+def load_plm(model_path='bert-base-uncased'):
+
+    tokenizer = AutoTokenizer.from_pretrained(model_path,)
+
+    print("Load Model:", model_path)
+
+    model = AutoModel.from_pretrained(model_path,low_cpu_mem_usage=True,)
+    return tokenizer, model
+
+def load_json(file):
+    with open(file, 'r') as f:
+        data = json.load(f)
+    return data
+
+def clean_text(raw_text):
+    if isinstance(raw_text, list):
+        new_raw_text=[]
+        for raw in raw_text:
+            raw = html.unescape(raw)
+            raw = re.sub(r'</?\w+[^>]*>', '', raw)
+            raw = re.sub(r'["\n\r]*', '', raw)
+            new_raw_text.append(raw.strip())
+        cleaned_text = ' '.join(new_raw_text)
+    else:
+        if isinstance(raw_text, dict):
+            cleaned_text = str(raw_text)[1:-1].strip()
+        else:
+            cleaned_text = raw_text.strip()
+        cleaned_text = html.unescape(cleaned_text)
+        cleaned_text = re.sub(r'</?\w+[^>]*>', '', cleaned_text)
+        cleaned_text = re.sub(r'["\n\r]*', '', cleaned_text)
+    index = -1
+    while -index < len(cleaned_text) and cleaned_text[index] == '.':
+        index -= 1
+    index += 1
+    if index == 0:
+        cleaned_text = cleaned_text + '.'
+    else:
+        cleaned_text = cleaned_text[:index] + '.'
+    if len(cleaned_text) >= 2000:
+        cleaned_text = ''
+    return cleaned_text
+
+def load_pickle(filename):
+    with open(filename, "rb") as f:
+        return pickle.load(f)
+
+
+def make_inters_in_order(inters):
+    user2inters, new_inters = collections.defaultdict(list), list()
+    for inter in inters:
+        user, item, rating, timestamp = inter
+        user2inters[user].append((user, item, rating, timestamp))
+    for user in user2inters:
+        user_inters = user2inters[user]
+        user_inters.sort(key=lambda d: d[3])
+        for inter in user_inters:
+            new_inters.append(inter)
+    return new_inters
+
+def write_json_file(dic, file):
+    print('Writing json file: ',file)
+    with open(file, 'w') as fp:
+        json.dump(dic, fp, indent=4)
+
+def write_remap_index(unit2index, file):
+    print('Writing remap file: ',file)
+    with open(file, 'w') as fp:
+        for unit in unit2index:
+            fp.write(unit + '\t' + str(unit2index[unit]) + '\n')
+
+
+intention_prompt = "After purchasing a {dataset_full_name} item named \"{item_title}\", the user left a comment expressing his opinion and personal preferences. The user's comment is as follows: \n\"{review}\" " \
+                    "\nAs we all know, user comments often contain information about both their personal preferences and the characteristics of the item they interacted with. From this comment, you can infer both the user's personal preferences and the characteristics of the item. " \
+                    "Please describe your inferred user preferences and item characteristics in the first person and in the following format:\n\nMy preferences: []\nThe item's characteristics: []\n\n" \
+                    "Note that your inference of the personalized preferences should not include any information about the title of the item."
+
+
+preference_prompt_1 = "Suppose the user has bought a variety of {dataset_full_name} items, they are: \n{item_titles}. \nAs we all know, these historically purchased items serve as a reflection of the user's personalized preferences. " \
+                        "Please analyze the user's personalized preferences based on the items he has bought and provide a brief third-person summary of the user's preferences, highlighting the key factors that influence his choice of items. Avoid listing specific items and do not list multiple examples. " \
+                        "Your analysis should be brief and in the third person."
+
+preference_prompt_2 = "Given a chronological list of {dataset_full_name} items that a user has purchased, we can analyze his long-term and short-term preferences. Long-term preferences are inherent characteristics of the user, which are reflected in all the items he has interacted with over time. Short-term preferences are the user's recent preferences, which are reflected in some of the items he has bought more recently. " \
+                        "To determine the user's long-term preferences, please analyze the contents of all the items he has bought. Look for common features that appear frequently across the user's shopping records. To determine the user's short-term preferences, focus on the items he has bought most recently. Identify any new or different features that have emerged in the user's shopping records. " \
+                        "Here is a chronological list of items that the user has bought: \n{item_titles}. \nPlease provide separate analyses for the user's long-term and short-term preferences. Your answer should be concise and general, without listing specific items. Your answer should be in the third person and in the following format:\n\nLong-term preferences: []\nShort-term preferences: []\n\n"
+
+
+# remove 'Magazine', 'Gift', 'Music', 'Kindle'
+amazon18_dataset_list = [
+    'Appliances', 'Beauty',
+    'Fashion', 'Software', 'Luxury', 'Scientific',  'Pantry',
+    'Instruments', 'Arts', 'Games', 'Office', 'Garden',
+    'Food', 'Cell', 'CDs', 'Automotive', 'Toys',
+    'Pet', 'Tools', 'Kindle', 'Sports', 'Movies',
+    'Electronics', 'Home', 'Clothing', 'Books'
+]
+
+amazon18_dataset2fullname = {
+    'Beauty': 'All_Beauty',
+    'Fashion': 'AMAZON_FASHION',
+    'Appliances': 'Appliances',
+    'Arts': 'Arts_Crafts_and_Sewing',
+    'Automotive': 'Automotive',
+    'Books': 'Books',
+    'CDs': 'CDs_and_Vinyl',
+    'Cell': 'Cell_Phones_and_Accessories',
+    'Clothing': 'Clothing_Shoes_and_Jewelry',
+    'Music': 'Digital_Music',
+    'Electronics': 'Electronics',
+    'Gift': 'Gift_Cards',
+    'Food': 'Grocery_and_Gourmet_Food',
+    'Home': 'Home_and_Kitchen',
+    'Scientific': 'Industrial_and_Scientific',
+    'Kindle': 'Kindle_Store',
+    'Luxury': 'Luxury_Beauty',
+    'Magazine': 'Magazine_Subscriptions',
+    'Movies': 'Movies_and_TV',
+    'Instruments': 'Musical_Instruments',
+    'Office': 'Office_Products',
+    'Garden': 'Patio_Lawn_and_Garden',
+    'Pet': 'Pet_Supplies',
+    'Pantry': 'Prime_Pantry',
+    'Software': 'Software',
+    'Sports': 'Sports_and_Outdoors',
+    'Tools': 'Tools_and_Home_Improvement',
+    'Toys': 'Toys_and_Games',
+    'Games': 'Video_Games'
+}
+
+amazon14_dataset_list = [
+    'Beauty','Toys','Sports'
+]
+
+amazon14_dataset2fullname = {
+    'Beauty': 'Beauty',
+    'Sports': 'Sports_and_Outdoors',
+    'Toys': 'Toys_and_Games',
+}
+
+# c1. c2. c3. c4.
+amazon_text_feature1 = ['title', 'category', 'brand']
+
+# re-order
+amazon_text_feature1_ro1 = ['brand', 'main_cat', 'category', 'title']
+
+# remove
+amazon_text_feature1_re1 = ['title']
+
+amazon_text_feature2 = ['title']
+
+amazon_text_feature3 = ['description']
+
+amazon_text_feature4 = ['description', 'main_cat', 'category', 'brand']
+
+amazon_text_feature5 = ['title', 'description']
+
+

evaluate.py

@@ -0,0 +1,66 @@
+import math
+
+def get_topk_results(predictions, scores, targets, k, all_items=None):
+    results = []
+    B = len(targets)
+    predictions = [_.split("Response:")[-1] for _ in predictions]
+    predictions = [_.strip().replace(" ","") for _ in predictions]
+
+    if all_items is not None:
+        for i, seq in enumerate(predictions):
+            if seq not in all_items:
+                scores[i] = -1000
+
+    for b in range(B):
+        batch_seqs = predictions[b * k: (b + 1) * k]
+        batch_scores = scores[b * k: (b + 1) * k]
+
+        pairs = [(a, b) for a, b in zip(batch_seqs, batch_scores)]
+        sorted_pairs = sorted(pairs, key=lambda x: x[1], reverse=True)
+        target_item = targets[b]
+        one_results = []
+        for sorted_pred in sorted_pairs:
+            if sorted_pred[0] == target_item:
+                one_results.append(1)
+            else:
+                one_results.append(0)
+
+        results.append(one_results)
+
+    return results
+
+def get_metrics_results(topk_results, metrics):
+    res = {}
+    for m in metrics:
+        if m.lower().startswith("hit"):
+            k = int(m.split("@")[1])
+            res[m] = hit_k(topk_results, k)
+        elif m.lower().startswith("ndcg"):
+            k = int(m.split("@")[1])
+            res[m] = ndcg_k(topk_results, k)
+        else:
+            raise NotImplementedError
+
+    return res
+
+
+def ndcg_k(topk_results, k):
+
+    ndcg = 0.0
+    for row in topk_results:
+        res = row[:k]
+        one_ndcg = 0.0
+        for i in range(len(res)):
+            one_ndcg += res[i] / math.log(i + 2, 2)
+        ndcg += one_ndcg
+    return ndcg
+
+
+def hit_k(topk_results, k):
+    hit = 0.0
+    for row in topk_results:
+        res = row[:k]
+        if sum(res) > 0:
+            hit += 1
+    return hit
+

finetune.py

@@ -0,0 +1,121 @@
+import argparse
+import os
+
+import sys
+from typing import List
+
+import torch
+import transformers
+
+from transformers import LlamaForCausalLM, LlamaTokenizer, LlamaConfig
+
+from utils import *
+from collator import Collator
+
+def train(args):
+
+    set_seed(args.seed)
+    ensure_dir(args.output_dir)
+
+    device_map = "auto"
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    ddp = world_size != 1
+    local_rank = int(os.environ.get("LOCAL_RANK") or 0)
+    if local_rank == 0:
+        print(vars(args))
+
+    if ddp:
+        device_map = {"": local_rank}
+
+    config = LlamaConfig.from_pretrained(args.base_model)
+    tokenizer = LlamaTokenizer.from_pretrained(
+        args.base_model,
+        model_max_length = args.model_max_length,
+        padding_side="right",
+    )
+    tokenizer.pad_token_id = 0
+    gradient_checkpointing = True
+
+    train_data, valid_data = load_datasets(args)
+    add_num = tokenizer.add_tokens(train_data.datasets[0].get_new_tokens())
+    config.vocab_size = len(tokenizer)
+    if local_rank == 0:
+        print("add {} new token.".format(add_num))
+        print("data num:", len(train_data))
+        tokenizer.save_pretrained(args.output_dir)
+        config.save_pretrained(args.output_dir)
+
+    collator = Collator(args, tokenizer)
+
+
+    model = LlamaForCausalLM.from_pretrained(
+        args.base_model,
+        # torch_dtype=torch.float16,
+        device_map=device_map,
+    )
+    model.resize_token_embeddings(len(tokenizer))
+
+
+    if not ddp and torch.cuda.device_count() > 1:
+        model.is_parallelizable = True
+        model.model_parallel = True
+
+
+    trainer = transformers.Trainer(
+        model=model,
+        train_dataset=train_data,
+        eval_dataset=valid_data,
+        args=transformers.TrainingArguments(
+            seed=args.seed,
+            per_device_train_batch_size=args.per_device_batch_size,
+            per_device_eval_batch_size=args.per_device_batch_size,
+            gradient_accumulation_steps=args.gradient_accumulation_steps,
+            warmup_ratio=args.warmup_ratio,
+            num_train_epochs=args.epochs,
+            learning_rate=args.learning_rate,
+            weight_decay=args.weight_decay,
+            lr_scheduler_type=args.lr_scheduler_type,
+            fp16=args.fp16,
+            bf16=args.bf16,
+            logging_steps=args.logging_step,
+            optim=args.optim,
+            gradient_checkpointing=gradient_checkpointing,
+            evaluation_strategy=args.save_and_eval_strategy,
+            save_strategy=args.save_and_eval_strategy,
+            eval_steps=args.save_and_eval_steps,
+            save_steps=args.save_and_eval_steps,
+            output_dir=args.output_dir,
+            save_total_limit=5,
+            load_best_model_at_end=True,
+            deepspeed=args.deepspeed,
+            ddp_find_unused_parameters=False if ddp else None,
+            report_to=None,
+            eval_delay= 1 if args.save_and_eval_strategy=="epoch" else 2000,
+            dataloader_num_workers = args.dataloader_num_workers,
+            dataloader_prefetch_factor = args.dataloader_prefetch_factor
+        ),
+        tokenizer=tokenizer,
+        data_collator=collator,
+    )
+    model.config.use_cache = False
+
+
+    trainer.train(
+        resume_from_checkpoint=args.resume_from_checkpoint,
+    )
+
+    trainer.save_state()
+    trainer.save_model(output_dir=args.output_dir)
+
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='LLMRec')
+    parser = parse_global_args(parser)
+    parser = parse_train_args(parser)
+    parser = parse_dataset_args(parser)
+
+    args = parser.parse_args()
+
+    train(args)

index/datasets.py

@@ -0,0 +1,21 @@
+import numpy as np
+import torch
+import torch.utils.data as data
+
+
+class EmbDataset(data.Dataset):
+
+    def __init__(self,data_path):
+
+        self.data_path = data_path
+        # self.embeddings = np.fromfile(data_path, dtype=np.float32).reshape(16859,-1)
+        self.embeddings = np.load(data_path)
+        self.dim = self.embeddings.shape[-1]
+
+    def __getitem__(self, index):
+        emb = self.embeddings[index]
+        tensor_emb=torch.FloatTensor(emb)
+        return tensor_emb
+
+    def __len__(self):
+        return len(self.embeddings)

index/generate_indices.py

@@ -0,0 +1,155 @@
+import collections
+import json
+import logging
+import argparse
+import numpy as np
+import torch
+from time import time
+from torch import optim
+from tqdm import tqdm
+
+from torch.utils.data import DataLoader
+
+from datasets import EmbDataset
+from models.rqvae import RQVAE
+
+import os
+
+def check_collision(all_indices_str):
+    tot_item = len(all_indices_str)
+    tot_indice = len(set(all_indices_str.tolist()))
+    return tot_item==tot_indice
+
+def get_indices_count(all_indices_str):
+    indices_count = collections.defaultdict(int)
+    for index in all_indices_str:
+        indices_count[index] += 1
+    return indices_count
+
+def get_collision_item(all_indices_str):
+    index2id = {}
+    for i, index in enumerate(all_indices_str):
+        if index not in index2id:
+            index2id[index] = []
+        index2id[index].append(i)
+
+    collision_item_groups = []
+
+    for index in index2id:
+        if len(index2id[index]) > 1:
+            collision_item_groups.append(index2id[index])
+
+    return collision_item_groups
+
+def parse_args():
+    parser = argparse.ArgumentParser(description = "Index")
+    parser.add_argument("--ckpt_path", type = str, default = "", help = "")
+    parser.add_argument("--data_path", type = str, default = "", help = "")
+    parser.add_argument("--save_path", type = str, default = "", help = "")
+    parser.add_argument("--device", type = str, default = "cuda:0", help = "gpu or cpu")
+    return parser.parse_args()
+
+infer_args = parse_args()
+print(infer_args)
+
+# dataset = "Games"
+# ckpt_path = "/zhengbowen/rqvae_ckpt/xxxx"
+# output_dir = f"/zhengbowen/data/{dataset}/"
+# output_file = f"{dataset}.index.json"
+# output_file = os.path.join(output_dir,output_file)
+# device = torch.device("cuda:1")
+
+device = torch.device(infer_args.device)
+
+ckpt = torch.load(infer_args.ckpt_path, map_location = torch.device('cpu'))
+args = ckpt["args"]
+state_dict = ckpt["state_dict"]
+
+data = EmbDataset(infer_args.data_path)
+
+model = RQVAE(in_dim=data.dim,
+                  num_emb_list=args.num_emb_list,
+                  e_dim=args.e_dim,
+                  layers=args.layers,
+                  dropout_prob=args.dropout_prob,
+                  bn=args.bn,
+                  loss_type=args.loss_type,
+                  quant_loss_weight=args.quant_loss_weight,
+                  kmeans_init=args.kmeans_init,
+                  kmeans_iters=args.kmeans_iters,
+                  sk_epsilons=args.sk_epsilons,
+                  sk_iters=args.sk_iters,
+                  )
+
+model.load_state_dict(state_dict)
+model = model.to(device)
+model.eval()
+print(model)
+
+data_loader = DataLoader(data,num_workers=args.num_workers,
+                             batch_size=64, shuffle=False,
+                             pin_memory=True)
+
+all_indices = []
+all_indices_str = []
+prefix = ["<a_{}>","<b_{}>","<c_{}>","<d_{}>","<e_{}>"]
+
+for d in tqdm(data_loader):
+    d = d.to(device)
+    indices = model.get_indices(d,use_sk=False)
+    indices = indices.view(-1, indices.shape[-1]).cpu().numpy()
+    for index in indices:
+        code = []
+        for i, ind in enumerate(index):
+            code.append(prefix[i].format(int(ind)))
+
+        all_indices.append(code)
+        all_indices_str.append(str(code))
+    # break
+
+all_indices = np.array(all_indices)
+all_indices_str = np.array(all_indices_str)
+
+for vq in model.rq.vq_layers[:-1]:
+    vq.sk_epsilon=0.0
+
+if model.rq.vq_layers[-1].sk_epsilon == 0.0:
+    model.rq.vq_layers[-1].sk_epsilon = 0.003
+
+tt = 0
+#There are often duplicate items in the dataset, and we no longer differentiate them
+while True:
+    if tt >= 10 or check_collision(all_indices_str):
+        break
+
+    collision_item_groups = get_collision_item(all_indices_str)
+    # print(collision_item_groups)
+    print(len(collision_item_groups))
+    for collision_items in collision_item_groups:
+        d = data[collision_items].to(device)
+
+        indices = model.get_indices(d, use_sk=True)
+        indices = indices.view(-1, indices.shape[-1]).cpu().numpy()
+        for item, index in zip(collision_items, indices):
+            code = []
+            for i, ind in enumerate(index):
+                code.append(prefix[i].format(int(ind)))
+
+            all_indices[item] = code
+            all_indices_str[item] = str(code)
+    tt += 1
+
+
+print("All indices number: ",len(all_indices))
+print("Max number of conflicts: ", max(get_indices_count(all_indices_str).values()))
+
+tot_item = len(all_indices_str)
+tot_indice = len(set(all_indices_str.tolist()))
+print("Collision Rate",(tot_item-tot_indice)/tot_item)
+
+all_indices_dict = {}
+for item, indices in enumerate(all_indices.tolist()):
+    all_indices_dict[item] = list(indices)
+
+with open(infer_args.save_path, 'w') as fp:
+    json.dump(all_indices_dict, fp)

index/main.py

@@ -0,0 +1,87 @@
+import argparse
+import random
+import torch
+import numpy as np
+from time import time
+import logging
+
+from torch.utils.data import DataLoader
+
+from datasets import EmbDataset
+from models.rqvae import RQVAE
+from trainer import  Trainer
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Index")
+
+    parser.add_argument('--lr', type=float, default=1e-3, help='learning rate')
+    parser.add_argument('--epochs', type=int, default=5000, help='number of epochs')
+    parser.add_argument('--batch_size', type=int, default=1024, help='batch size')
+    parser.add_argument('--num_workers', type=int, default=4, )
+    parser.add_argument('--eval_step', type=int, default=50, help='eval step')
+    parser.add_argument('--learner', type=str, default="AdamW", help='optimizer')
+    parser.add_argument("--data_path", type=str,
+                        default="../data/Games/Games.emb-llama-td.npy",
+                        help="Input data path.")
+
+    parser.add_argument('--weight_decay', type=float, default=1e-4, help='l2 regularization weight')
+    parser.add_argument("--dropout_prob", type=float, default=0.0, help="dropout ratio")
+    parser.add_argument("--bn", type=bool, default=False, help="use bn or not")
+    parser.add_argument("--loss_type", type=str, default="mse", help="loss_type")
+    parser.add_argument("--kmeans_init", type=bool, default=True, help="use kmeans_init or not")
+    parser.add_argument("--kmeans_iters", type=int, default=100, help="max kmeans iters")
+    parser.add_argument('--sk_epsilons', type=float, nargs='+', default=[0.0, 0.0, 0.0], help="sinkhorn epsilons")
+    parser.add_argument("--sk_iters", type=int, default=50, help="max sinkhorn iters")
+
+    parser.add_argument("--device", type=str, default="cuda:1", help="gpu or cpu")
+
+    parser.add_argument('--num_emb_list', type=int, nargs='+', default=[256,256,256], help='emb num of every vq')
+    parser.add_argument('--e_dim', type=int, default=32, help='vq codebook embedding size')
+    parser.add_argument('--quant_loss_weight', type=float, default=1.0, help='vq quantion loss weight')
+    parser.add_argument('--layers', type=int, nargs='+', default=[2048,1024,512,256,128,64], help='hidden sizes of every layer')
+
+    parser.add_argument("--ckpt_dir", type=str, default="", help="output directory for model")
+
+    return parser.parse_args()
+
+
+if __name__ == '__main__':
+    """fix the random seed"""
+    seed = 2023
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+    args = parse_args()
+    print(args)
+
+    logging.basicConfig(level=logging.DEBUG)
+
+    """build dataset"""
+    data = EmbDataset(args.data_path)
+    model = RQVAE(in_dim=data.dim,
+                  num_emb_list=args.num_emb_list,
+                  e_dim=args.e_dim,
+                  layers=args.layers,
+                  dropout_prob=args.dropout_prob,
+                  bn=args.bn,
+                  loss_type=args.loss_type,
+                  quant_loss_weight=args.quant_loss_weight,
+                  kmeans_init=args.kmeans_init,
+                  kmeans_iters=args.kmeans_iters,
+                  sk_epsilons=args.sk_epsilons,
+                  sk_iters=args.sk_iters,
+                  )
+    print(model)
+    data_loader = DataLoader(data,num_workers=args.num_workers,
+                             batch_size=args.batch_size, shuffle=True,
+                             pin_memory=True)
+    trainer = Trainer(args,model)
+    best_loss, best_collision_rate = trainer.fit(data_loader)
+
+    print("Best Loss",best_loss)
+    print("Best Collision Rate", best_collision_rate)
+

index/models/layers.py

@@ -0,0 +1,106 @@
+import torch
+import torch.nn as nn
+from torch.nn.init import xavier_normal_
+from sklearn.cluster import KMeans
+
+
+class MLPLayers(nn.Module):
+
+    def __init__(
+        self, layers, dropout=0.0, activation="relu", bn=False
+    ):
+        super(MLPLayers, self).__init__()
+        self.layers = layers
+        self.dropout = dropout
+        self.activation = activation
+        self.use_bn = bn
+
+        mlp_modules = []
+        for idx, (input_size, output_size) in enumerate(
+            zip(self.layers[:-1], self.layers[1:])
+        ):
+            mlp_modules.append(nn.Dropout(p=self.dropout))
+            mlp_modules.append(nn.Linear(input_size, output_size))
+            if self.use_bn:
+                mlp_modules.append(nn.BatchNorm1d(num_features=output_size))
+            activation_func = activation_layer(self.activation, output_size)
+            if activation_func is not None and idx != (len(self.layers)-2):
+                mlp_modules.append(activation_func)
+
+        self.mlp_layers = nn.Sequential(*mlp_modules)
+        self.apply(self.init_weights)
+
+    def init_weights(self, module):
+        # We just initialize the module with normal distribution as the paper said
+        if isinstance(module, nn.Linear):
+            xavier_normal_(module.weight.data)
+            if module.bias is not None:
+                module.bias.data.fill_(0.0)
+
+    def forward(self, input_feature):
+        return self.mlp_layers(input_feature)
+
+def activation_layer(activation_name="relu", emb_dim=None):
+
+    if activation_name is None:
+        activation = None
+    elif isinstance(activation_name, str):
+        if activation_name.lower() == "sigmoid":
+            activation = nn.Sigmoid()
+        elif activation_name.lower() == "tanh":
+            activation = nn.Tanh()
+        elif activation_name.lower() == "relu":
+            activation = nn.ReLU()
+        elif activation_name.lower() == "leakyrelu":
+            activation = nn.LeakyReLU()
+        elif activation_name.lower() == "none":
+            activation = None
+    elif issubclass(activation_name, nn.Module):
+        activation = activation_name()
+    else:
+        raise NotImplementedError(
+            "activation function {} is not implemented".format(activation_name)
+        )
+
+    return activation
+
+def kmeans(
+    samples,
+    num_clusters,
+    num_iters = 10,
+):
+    B, dim, dtype, device = samples.shape[0], samples.shape[-1], samples.dtype, samples.device
+    x = samples.cpu().detach().numpy()
+
+    cluster = KMeans(n_clusters = num_clusters, max_iter = num_iters).fit(x)
+
+    centers = cluster.cluster_centers_
+    tensor_centers = torch.from_numpy(centers).to(device)
+
+    return tensor_centers
+
+
+@torch.no_grad()
+def sinkhorn_algorithm(distances, epsilon, sinkhorn_iterations):
+    Q = torch.exp(- distances / epsilon)
+
+    B = Q.shape[0] # number of samples to assign
+    K = Q.shape[1] # how many centroids per block (usually set to 256)
+
+    # make the matrix sums to 1
+    sum_Q = Q.sum(-1, keepdim=True).sum(-2, keepdim=True)
+    Q /= sum_Q
+    # print(Q.sum())
+    for it in range(sinkhorn_iterations):
+
+        # normalize each column: total weight per sample must be 1/B
+        Q /= torch.sum(Q, dim=1, keepdim=True)
+        Q /= B
+
+        # normalize each row: total weight per prototype must be 1/K
+        Q /= torch.sum(Q, dim=0, keepdim=True)
+        Q /= K
+
+
+    Q *= B # the colomns must sum to 1 so that Q is an assignment
+    return Q

index/models/rq.py

@@ -0,0 +1,54 @@
+import torch
+import torch.nn as nn
+
+from .vq import VectorQuantizer
+
+
+class ResidualVectorQuantizer(nn.Module):
+    """ References:
+        SoundStream: An End-to-End Neural Audio Codec
+        https://arxiv.org/pdf/2107.03312.pdf
+    """
+
+    def __init__(self, n_e_list, e_dim, sk_epsilons,
+                 kmeans_init = False, kmeans_iters = 100, sk_iters=100,):
+        super().__init__()
+        self.n_e_list = n_e_list
+        self.e_dim = e_dim
+        self.num_quantizers = len(n_e_list)
+        self.kmeans_init = kmeans_init
+        self.kmeans_iters = kmeans_iters
+        self.sk_epsilons = sk_epsilons
+        self.sk_iters = sk_iters
+        self.vq_layers = nn.ModuleList([VectorQuantizer(n_e, e_dim,
+                                                        kmeans_init = self.kmeans_init,
+                                                        kmeans_iters = self.kmeans_iters,
+                                                        sk_epsilon=sk_epsilon,
+                                                        sk_iters=sk_iters)
+                                        for n_e, sk_epsilon in zip(n_e_list,sk_epsilons) ])
+
+    def get_codebook(self):
+        all_codebook = []
+        for quantizer in self.vq_layers:
+            codebook = quantizer.get_codebook()
+            all_codebook.append(codebook)
+        return torch.stack(all_codebook)
+
+    def forward(self, x, use_sk=True):
+        all_losses = []
+        all_indices = []
+
+        x_q = 0
+        residual = x
+        for quantizer in self.vq_layers:
+            x_res, loss, indices = quantizer(residual, use_sk=use_sk)
+            residual = residual - x_res
+            x_q = x_q + x_res
+
+            all_losses.append(loss)
+            all_indices.append(indices)
+
+        mean_losses = torch.stack(all_losses).mean()
+        all_indices = torch.stack(all_indices, dim=-1)
+
+        return x_q, mean_losses, all_indices

index/models/rqvae.py

@@ -0,0 +1,82 @@
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from .layers import MLPLayers
+from .rq import ResidualVectorQuantizer
+
+
+class RQVAE(nn.Module):
+    def __init__(self,
+                 in_dim=768,
+                 # num_emb_list=[256,256,256,256],
+                 num_emb_list=None,
+                 e_dim=64,
+                 # layers=[512,256,128],
+                 layers=None,
+                 dropout_prob=0.0,
+                 bn=False,
+                 loss_type="mse",
+                 quant_loss_weight=1.0,
+                 kmeans_init=False,
+                 kmeans_iters=100,
+                 # sk_epsilons=[0,0,0.003,0.01]],
+                 sk_epsilons=None,
+                 sk_iters=100,
+        ):
+        super(RQVAE, self).__init__()
+
+        self.in_dim = in_dim
+        self.num_emb_list = num_emb_list
+        self.e_dim = e_dim
+
+        self.layers = layers
+        self.dropout_prob = dropout_prob
+        self.bn = bn
+        self.loss_type = loss_type
+        self.quant_loss_weight=quant_loss_weight
+        self.kmeans_init = kmeans_init
+        self.kmeans_iters = kmeans_iters
+        self.sk_epsilons = sk_epsilons
+        self.sk_iters = sk_iters
+
+        self.encode_layer_dims = [self.in_dim] + self.layers + [self.e_dim]
+        self.encoder = MLPLayers(layers=self.encode_layer_dims,
+                                 dropout=self.dropout_prob,bn=self.bn)
+
+        self.rq = ResidualVectorQuantizer(num_emb_list, e_dim,
+                                          kmeans_init = self.kmeans_init,
+                                          kmeans_iters = self.kmeans_iters,
+                                          sk_epsilons=self.sk_epsilons,
+                                          sk_iters=self.sk_iters,)
+
+        self.decode_layer_dims = self.encode_layer_dims[::-1]
+        self.decoder = MLPLayers(layers=self.decode_layer_dims,
+                                       dropout=self.dropout_prob,bn=self.bn)
+
+    def forward(self, x, use_sk=True):
+        x = self.encoder(x)
+        x_q, rq_loss, indices = self.rq(x,use_sk=use_sk)
+        out = self.decoder(x_q)
+
+        return out, rq_loss, indices
+
+    @torch.no_grad()
+    def get_indices(self, xs, use_sk=False):
+        x_e = self.encoder(xs)
+        _, _, indices = self.rq(x_e, use_sk=use_sk)
+        return indices
+
+    def compute_loss(self, out, quant_loss, xs=None):
+
+        if self.loss_type == 'mse':
+            loss_recon = F.mse_loss(out, xs, reduction='mean')
+        elif self.loss_type == 'l1':
+            loss_recon = F.l1_loss(out, xs, reduction='mean')
+        else:
+            raise ValueError('incompatible loss type')
+
+        loss_total = loss_recon + self.quant_loss_weight * quant_loss
+
+        return loss_total, loss_recon

index/models/vq.py

@@ -0,0 +1,103 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .layers import kmeans, sinkhorn_algorithm
+
+
+class VectorQuantizer(nn.Module):
+
+    def __init__(self, n_e, e_dim,
+                 beta = 0.25, kmeans_init = False, kmeans_iters = 10,
+                 sk_epsilon=0.01, sk_iters=100):
+        super().__init__()
+        self.n_e = n_e
+        self.e_dim = e_dim
+        self.beta = beta
+        self.kmeans_init = kmeans_init
+        self.kmeans_iters = kmeans_iters
+        self.sk_epsilon = sk_epsilon
+        self.sk_iters = sk_iters
+
+        self.embedding = nn.Embedding(self.n_e, self.e_dim)
+        if not kmeans_init:
+            self.initted = True
+            self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e)
+        else:
+            self.initted = False
+            self.embedding.weight.data.zero_()
+
+    def get_codebook(self):
+        return self.embedding.weight
+
+    def get_codebook_entry(self, indices, shape=None):
+        # get quantized latent vectors
+        z_q = self.embedding(indices)
+        if shape is not None:
+            z_q = z_q.view(shape)
+
+        return z_q
+
+    def init_emb(self, data):
+
+        centers = kmeans(
+            data,
+            self.n_e,
+            self.kmeans_iters,
+        )
+
+        self.embedding.weight.data.copy_(centers)
+        self.initted = True
+
+    @staticmethod
+    def center_distance_for_constraint(distances):
+        # distances: B, K
+        max_distance = distances.max()
+        min_distance = distances.min()
+
+        middle = (max_distance + min_distance) / 2
+        amplitude = max_distance - middle + 1e-5
+        assert amplitude > 0
+        centered_distances = (distances - middle) / amplitude
+        return centered_distances
+
+    def forward(self, x, use_sk=True):
+        # Flatten input
+        latent = x.view(-1, self.e_dim)
+
+        if not self.initted and self.training:
+            self.init_emb(latent)
+
+        # Calculate the L2 Norm between latent and Embedded weights
+        d = torch.sum(latent**2, dim=1, keepdim=True) + \
+            torch.sum(self.embedding.weight**2, dim=1, keepdim=True).t()- \
+            2 * torch.matmul(latent, self.embedding.weight.t())
+        if not use_sk or self.sk_epsilon <= 0:
+            indices = torch.argmin(d, dim=-1)
+            # print("=======",self.sk_epsilon)
+        else:
+            # print("++++++++",self.sk_epsilon)
+            d = self.center_distance_for_constraint(d)
+            d = d.double()
+            Q = sinkhorn_algorithm(d,self.sk_epsilon,self.sk_iters)
+            # print(Q.sum(0)[:10])
+            if torch.isnan(Q).any() or torch.isinf(Q).any():
+                print(f"Sinkhorn Algorithm returns nan/inf values.")
+            indices = torch.argmax(Q, dim=-1)
+
+        # indices = torch.argmin(d, dim=-1)
+
+        x_q = self.embedding(indices).view(x.shape)
+
+        # compute loss for embedding
+        commitment_loss = F.mse_loss(x_q.detach(), x)
+        codebook_loss = F.mse_loss(x_q, x.detach())
+        loss = codebook_loss + self.beta * commitment_loss
+
+        # preserve gradients
+        x_q = x + (x_q - x).detach()
+
+        indices = indices.view(x.shape[:-1])
+
+        return x_q, loss, indices
+
+

index/run.sh

@@ -0,0 +1,18 @@
+DATA_PATH=/datain/v-yinju/rqvae-zzx/data/instruments-emb-llama.npy
+CKPT_DIR=your_ckpt_save_directory # E.g., /datain/v-yinju/rqvae-zzx/model
+mkdir -p $CKPT_DIR
+python -u main.py \
+  --num_emb_list 256 256 256 256 \
+  --sk_epsilons 0.0 0.0 0.0 0.003 \
+  --lr 1e-3 \
+  --device cuda:0 \
+  --batch_size 1024 \
+  --data_path $DATA_PATH \
+  --ckpt_dir $CKPT_DIR
+
+# Infer item index
+# python generate_indices.py \
+#  --ckpt_path your_rqvae_model_path \ E.g., /datain/v-yinju/rqvae-zzx/model/20241127/best_collision_model.pth
+#  --data_path $DATA_PATH \
+#  --save_path your_index_save_path \ E.g., /datain/v-yinju/rqvae-zzx/model/20241127/indices.json
+#  --device cuda:0

index/trainer.py

@@ -0,0 +1,209 @@
+import logging
+
+import numpy as np
+import torch
+from time import time
+from torch import optim
+from tqdm import tqdm
+
+from utils import ensure_dir,set_color,get_local_time
+import os
+
+class Trainer(object):
+
+    def __init__(self, args, model):
+        self.args = args
+        self.model = model
+        self.logger = logging.getLogger()
+
+        self.lr = args.lr
+        self.learner = args.learner
+        self.weight_decay = args.weight_decay
+        self.epochs = args.epochs
+        self.eval_step = min(args.eval_step, self.epochs)
+        self.device = args.device
+        self.device = torch.device(self.device)
+        self.ckpt_dir = args.ckpt_dir
+        saved_model_dir = "{}".format(get_local_time())
+        self.ckpt_dir = os.path.join(self.ckpt_dir,saved_model_dir)
+        ensure_dir(self.ckpt_dir)
+
+        self.best_loss = np.inf
+        self.best_collision_rate = np.inf
+        self.best_loss_ckpt = "best_loss_model.pth"
+        self.best_collision_ckpt = "best_collision_model.pth"
+        self.optimizer = self._build_optimizer()
+        self.model = self.model.to(self.device)
+
+    def _build_optimizer(self):
+
+        params = self.model.parameters()
+        learner =  self.learner
+        learning_rate = self.lr
+        weight_decay = self.weight_decay
+
+        if learner.lower() == "adam":
+            optimizer = optim.Adam(params, lr=learning_rate, weight_decay=weight_decay)
+        elif learner.lower() == "sgd":
+            optimizer = optim.SGD(params, lr=learning_rate, weight_decay=weight_decay)
+        elif learner.lower() == "adagrad":
+            optimizer = optim.Adagrad(
+                params, lr=learning_rate, weight_decay=weight_decay
+            )
+            for state in optimizer.state.values():
+                for k, v in state.items():
+                    if torch.is_tensor(v):
+                        state[k] = v.to(self.device)
+        elif learner.lower() == "rmsprop":
+            optimizer = optim.RMSprop(
+                params, lr=learning_rate, weight_decay=weight_decay
+            )
+        elif learner.lower() == 'adamw':
+            optimizer = optim.AdamW(
+                params, lr=learning_rate, weight_decay=weight_decay
+            )
+        else:
+            self.logger.warning(
+                "Received unrecognized optimizer, set default Adam optimizer"
+            )
+            optimizer = optim.Adam(params, lr=learning_rate)
+        return optimizer
+    def _check_nan(self, loss):
+        if torch.isnan(loss):
+            raise ValueError("Training loss is nan")
+
+    def _train_epoch(self, train_data, epoch_idx):
+
+        self.model.train()
+
+        total_loss = 0
+        total_recon_loss = 0
+        iter_data = tqdm(
+                    train_data,
+                    total=len(train_data),
+                    ncols=100,
+                    desc=set_color(f"Train {epoch_idx}","pink"),
+                    )
+
+        for batch_idx, data in enumerate(iter_data):
+            data = data.to(self.device)
+            self.optimizer.zero_grad()
+            out, rq_loss, indices = self.model(data)
+            loss, loss_recon = self.model.compute_loss(out, rq_loss, xs=data)
+            self._check_nan(loss)
+            loss.backward()
+            self.optimizer.step()
+            total_loss += loss.item()
+            total_recon_loss += loss_recon.item()
+
+        return total_loss, total_recon_loss
+
+    @torch.no_grad()
+    def _valid_epoch(self, valid_data):
+
+        self.model.eval()
+
+        iter_data =tqdm(
+                valid_data,
+                total=len(valid_data),
+                ncols=100,
+                desc=set_color(f"Evaluate   ", "pink"),
+            )
+        indices_set = set()
+        num_sample = 0
+        for batch_idx, data in enumerate(iter_data):
+            num_sample += len(data)
+            data = data.to(self.device)
+            indices = self.model.get_indices(data)
+            indices = indices.view(-1,indices.shape[-1]).cpu().numpy()
+            for index in indices:
+                code = "-".join([str(int(_)) for _ in index])
+                indices_set.add(code)
+
+        collision_rate = (num_sample - len(indices_set))/num_sample
+
+        return collision_rate
+
+    def _save_checkpoint(self, epoch, collision_rate=1, ckpt_file=None):
+
+        ckpt_path = os.path.join(self.ckpt_dir,ckpt_file) if ckpt_file \
+            else os.path.join(self.ckpt_dir, 'epoch_%d_collision_%.4f_model.pth' % (epoch, collision_rate))
+        state = {
+            "args": self.args,
+            "epoch": epoch,
+            "best_loss": self.best_loss,
+            "best_collision_rate": self.best_collision_rate,
+            "state_dict": self.model.state_dict(),
+            "optimizer": self.optimizer.state_dict(),
+        }
+        torch.save(state, ckpt_path, pickle_protocol=4)
+
+        self.logger.info(
+            set_color("Saving current", "blue") + f": {ckpt_path}"
+        )
+
+    def _generate_train_loss_output(self, epoch_idx, s_time, e_time, loss, recon_loss):
+        train_loss_output = (
+            set_color("epoch %d training", "green")
+            + " ["
+            + set_color("time", "blue")
+            + ": %.2fs, "
+        ) % (epoch_idx, e_time - s_time)
+        train_loss_output += set_color("train loss", "blue") + ": %.4f" % loss
+        train_loss_output +=", "
+        train_loss_output += set_color("reconstruction loss", "blue") + ": %.4f" % recon_loss
+        return train_loss_output + "]"
+
+
+    def fit(self, data):
+
+        cur_eval_step = 0
+
+        for epoch_idx in range(self.epochs):
+            # train
+            training_start_time = time()
+            train_loss, train_recon_loss = self._train_epoch(data, epoch_idx)
+            training_end_time = time()
+            train_loss_output = self._generate_train_loss_output(
+                epoch_idx, training_start_time, training_end_time, train_loss, train_recon_loss
+            )
+            self.logger.info(train_loss_output)
+
+            if train_loss < self.best_loss:
+                self.best_loss = train_loss
+                # self._save_checkpoint(epoch=epoch_idx,ckpt_file=self.best_loss_ckpt)
+
+            # eval
+            if (epoch_idx + 1) % self.eval_step == 0:
+                valid_start_time = time()
+                collision_rate = self._valid_epoch(data)
+
+                if collision_rate < self.best_collision_rate:
+                    self.best_collision_rate = collision_rate
+                    cur_eval_step = 0
+                    self._save_checkpoint(epoch_idx, collision_rate=collision_rate,
+                                          ckpt_file=self.best_collision_ckpt)
+                else:
+                    cur_eval_step += 1
+
+
+                valid_end_time = time()
+                valid_score_output = (
+                    set_color("epoch %d evaluating", "green")
+                    + " ["
+                    + set_color("time", "blue")
+                    + ": %.2fs, "
+                    + set_color("collision_rate", "blue")
+                    + ": %f]"
+                ) % (epoch_idx, valid_end_time - valid_start_time, collision_rate)
+
+                self.logger.info(valid_score_output)
+                if epoch_idx>1000:
+                    self._save_checkpoint(epoch_idx, collision_rate=collision_rate)
+
+
+        return self.best_loss, self.best_collision_rate
+
+
+
+

index/utils.py

@@ -0,0 +1,36 @@
+
+import datetime
+import os
+
+
+def ensure_dir(dir_path):
+
+    os.makedirs(dir_path, exist_ok=True)
+
+def set_color(log, color, highlight=True):
+    color_set = ["black", "red", "green", "yellow", "blue", "pink", "cyan", "white"]
+    try:
+        index = color_set.index(color)
+    except:
+        index = len(color_set) - 1
+    prev_log = "\033["
+    if highlight:
+        prev_log += "1;3"
+    else:
+        prev_log += "0;3"
+    prev_log += str(index) + "m"
+    return prev_log + log + "\033[0m"
+
+def get_local_time():
+    r"""Get current time
+
+    Returns:
+        str: current time
+    """
+    cur = datetime.datetime.now()
+    cur = cur.strftime("%b-%d-%Y_%H-%M-%S")
+
+    return cur
+
+
+

instruments_eval.sh

@@ -0,0 +1,17 @@
+DATASET=Instruments
+BASE=/datain/v-yinju/llama-7b
+DATA_PATH=/datain/v-yinju/rqvae-zzx/data
+CKPT_PATH=/datain/v-yinju/RQVAE_Bench/llama
+RESULTS_FILE=$CKPT_PATH/result.json
+INDEX=/datain/v-yinju/RQVAE_Bench/rqvae/Nov-27-2024_23-08-08/indices.json
+
+torchrun --nproc_per_node=8 test_ddp.py \
+    --base_model $BASE \
+    --ckpt_path $CKPT_PATH \
+    --dataset $DATASET \
+    --data_path $DATA_PATH \
+    --results_file $RESULTS_FILE \
+    --test_batch_size 1 \
+    --num_beams 10 \
+    --test_prompt_ids all \
+    --index_file $INDEX

instruments_train.sh

@@ -0,0 +1,34 @@
+export WANDB_MODE=disabled
+export CUDA_LAUNCH_BLOCKING=1
+
+DATASET=Instruments
+BASE_MODEL=/datain/v-yinju/llama-7b
+DATA_PATH=/datain/v-yinju/rqvae-zzx/data
+INDEX=your_index_save_path
+OUTPUT_DIR=your_ckpt_save_dir
+
+mkdir -p $OUTPUT_DIR
+
+torchrun --nproc_per_node=8 lora_finetune.py \
+    --base_model $BASE_MODEL \
+    --output_dir $OUTPUT_DIR \
+    --dataset $DATASET \
+    --data_path $DATA_PATH \
+    --per_device_batch_size 6 \
+    --gradient_accumulation_steps 2 \
+    --learning_rate 5e-5 \
+    --epochs 4 \
+    --weight_decay 0.01 \
+    --save_and_eval_strategy epoch \
+    --fp16 \
+    --deepspeed ./config/ds_z2_fp16.json \
+    --dataloader_num_workers 4 \
+    --only_train_response \
+    --tasks seqrec,item2index,index2item,fusionseqrec,itemsearch,preferenceobtain \
+    --train_prompt_sample_num 1,1,1,1,1,1 \
+    --train_data_sample_num 0,0,0,0,0,0 \
+    --index_file $INDEX
+
+cd convert
+nohup ./convert.sh $OUTPUT_DIR >convert.log 2>&1 &
+cd ..

lora_finetune.py

@@ -0,0 +1,164 @@
+import argparse
+import os
+import sys
+from typing import List
+
+import torch
+import transformers
+
+
+from peft import (
+    TaskType,
+    LoraConfig,
+    get_peft_model,
+    get_peft_model_state_dict,
+    set_peft_model_state_dict,
+)
+from transformers import LlamaForCausalLM, LlamaTokenizer, LlamaConfig
+
+from utils import *
+from collator import Collator
+
+def train(args):
+
+    set_seed(args.seed)
+    ensure_dir(args.output_dir)
+
+    device_map = "auto"
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    ddp = world_size != 1
+    local_rank = int(os.environ.get("LOCAL_RANK") or 0)
+    if local_rank == 0:
+        print(vars(args))
+
+    if ddp:
+        device_map = {"": local_rank}
+
+    config = LlamaConfig.from_pretrained(args.base_model)
+    tokenizer = LlamaTokenizer.from_pretrained(
+        args.base_model,
+        model_max_length=args.model_max_length,
+        padding_side="right",
+    )
+    tokenizer.pad_token_id = 0
+
+    train_data, valid_data = load_datasets(args)
+    add_num = tokenizer.add_tokens(train_data.datasets[0].get_new_tokens())
+    config.vocab_size = len(tokenizer)
+    if local_rank == 0:
+        print("add {} new token.".format(add_num))
+        print("data num:", len(train_data))
+        tokenizer.save_pretrained(args.output_dir)
+        config.save_pretrained(args.output_dir)
+
+    collator = Collator(args, tokenizer)
+
+    model = LlamaForCausalLM.from_pretrained(
+        args.base_model,
+        torch_dtype=torch.float16,
+        device_map=device_map,
+    )
+    model.resize_token_embeddings(len(tokenizer))
+
+    config = LoraConfig(
+        r=args.lora_r,
+        lora_alpha=args.lora_alpha,
+        target_modules=args.lora_target_modules.split(","),
+        modules_to_save=args.lora_modules_to_save.split(","),
+        lora_dropout=args.lora_dropout,
+        bias="none",
+        inference_mode=False,
+        task_type=TaskType.CAUSAL_LM,
+    )
+    model = get_peft_model(model, config)
+
+    if args.resume_from_checkpoint:
+        checkpoint_name = os.path.join(
+            args.resume_from_checkpoint, "adapter_model.bin"
+        )  # only LoRA model - LoRA config above has to fit
+        args.resume_from_checkpoint = False  # So the trainer won't try loading its state
+        # The two files above have a different name depending on how they were saved, but are actually the same.
+        if os.path.exists(checkpoint_name):
+            if local_rank == 0:
+                print(f"Restarting from {checkpoint_name}")
+            adapters_weights = torch.load(checkpoint_name)
+            model = set_peft_model_state_dict(model, adapters_weights)
+        else:
+            if local_rank == 0:
+                print(f"Checkpoint {checkpoint_name} not found")
+
+    for n, p in model.named_parameters():
+        if "original_module" in n and any(module_name in n for module_name in config.modules_to_save):
+            p.requires_grad = False
+
+    if local_rank == 0:
+        model.print_trainable_parameters()
+
+
+    if not ddp and torch.cuda.device_count() > 1:
+        model.is_parallelizable = True
+        model.model_parallel = True
+
+    trainer = transformers.Trainer(
+        model=model,
+        train_dataset=train_data,
+        eval_dataset=valid_data,
+        args=transformers.TrainingArguments(
+            seed=args.seed,
+            per_device_train_batch_size=args.per_device_batch_size,
+            per_device_eval_batch_size=args.per_device_batch_size,
+            gradient_accumulation_steps=args.gradient_accumulation_steps,
+            warmup_ratio=args.warmup_ratio,
+            num_train_epochs=args.epochs,
+            learning_rate=args.learning_rate,
+            weight_decay=args.weight_decay,
+            lr_scheduler_type=args.lr_scheduler_type,
+            fp16=args.fp16,
+            bf16=args.bf16,
+            logging_steps=args.logging_step,
+            optim=args.optim,
+            gradient_checkpointing=True,
+            evaluation_strategy=args.save_and_eval_strategy,
+            save_strategy=args.save_and_eval_strategy,
+            eval_steps=args.save_and_eval_steps,
+            save_steps=args.save_and_eval_steps,
+            output_dir=args.output_dir,
+            save_total_limit=5,
+            load_best_model_at_end=True,
+            deepspeed=args.deepspeed,
+            ddp_find_unused_parameters=False if ddp else None,
+            report_to=None,
+            eval_delay=1 if args.save_and_eval_strategy=="epoch" else 2000,
+            dataloader_num_workers = args.dataloader_num_workers,
+            dataloader_prefetch_factor = args.dataloader_prefetch_factor
+        ),
+        tokenizer=tokenizer,
+        data_collator=collator,
+    )
+    model.config.use_cache = False
+
+    # old_state_dict = model.state_dict
+    # model.state_dict = (
+    #     lambda self, *_, **__: get_peft_model_state_dict(self, old_state_dict())
+    # ).__get__(model, type(model))
+
+    if torch.__version__ >= "2" and sys.platform != "win32":
+        model = torch.compile(model)
+
+    trainer.train(
+        resume_from_checkpoint=args.resume_from_checkpoint,
+    )
+
+    trainer.save_state()
+    trainer.save_model(output_dir=args.output_dir)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='LLMRec')
+    parser = parse_global_args(parser)
+    parser = parse_train_args(parser)
+    parser = parse_dataset_args(parser)
+
+    args = parser.parse_args()
+
+    train(args)

prompt.py

@@ -0,0 +1,663 @@
+
+
+sft_prompt = "Below is an instruction that describes a task. Write a response that appropriately completes the request." \
+             "\n\n### Instruction:\n{instruction}\n\n### Response:{response}"
+
+
+
+
+
+
+
+all_prompt = {}
+
+# =====================================================
+# Task 1 -- Sequential Recommendation -- 17 Prompt
+# =====================================================
+
+seqrec_prompt = []
+
+#####——0
+prompt = {}
+prompt["instruction"] = "The user has interacted with items {inters} in chronological order. Can you predict the next possible item that the user may expect?"
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+#####——1
+prompt = {}
+prompt["instruction"] = "I find the user's historical interactive items: {inters}, and I want to know what next item the user needs. Can you help me decide?"
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+#####——2
+prompt = {}
+prompt["instruction"] = "Here are the user's historical interactions: {inters}, try to recommend another item to the user. Note that the historical interactions are arranged in chronological order."
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+#####——3
+prompt = {}
+prompt["instruction"] = "Based on the items that the user has interacted with: {inters}, can you determine what item would be recommended to him next?"
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+#####——4
+prompt = {}
+prompt["instruction"] = "The user has interacted with the following items in order: {inters}. What else do you think the user need?"
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+#####——5
+prompt = {}
+prompt["instruction"] = "Here is the item interaction history of the user: {inters}, what to recommend to the user next?"
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+#####——6
+prompt = {}
+prompt["instruction"] = "Which item would the user be likely to interact with next after interacting with items {inters}?"
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+#####——7
+prompt = {}
+prompt["instruction"] = "By analyzing the user's historical interactions with items {inters}, what is the next expected interaction item?"
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+#####——8
+prompt = {}
+prompt["instruction"] = "After interacting with items {inters}, what is the next item that could be recommended for the user?"
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+#####——9
+prompt = {}
+prompt["instruction"] = "Given the user's historical interactive items arranged in chronological order: {inters}, can you recommend a suitable item for the user?"
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+#####——10
+prompt = {}
+prompt["instruction"] = "Considering the user has interacted with items {inters}. What is the next recommendation for the user?"
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+#####——11
+prompt = {}
+prompt["instruction"] = "What is the top recommended item for the user who has previously interacted with items {inters} in order?"
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+#####——12
+prompt = {}
+prompt["instruction"] = "The user has interacted with the following items in the past in order: {inters}. Please predict the next item that the user most desires based on the given interaction records."
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+# prompt = {}
+# prompt["instruction"] = "The user has interacted with the following items in the past in order: {inters}. Please predict the next item that the user is most likely to interact with based on the given interaction record. Note that his most recently interacted item is {}."
+# prompt["response"] = "{item}"
+# prompt["task"] = "sequential"
+# prompt["id"] = "1-13"
+#
+# seqrec_prompt.append(prompt)
+
+#####——13
+prompt = {}
+prompt["instruction"] = "Using the user's historical interactions as input data, suggest the next item that the user is highly likely to enjoy. The historical interactions are provided as follows: {inters}."
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+#####——14
+prompt = {}
+prompt["instruction"] = "You can access the user's historical item interaction records: {inters}. Now your task is to recommend the next potential item to him, considering his past interactions."
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+#####——15
+prompt = {}
+prompt["instruction"] = "You have observed that the user has interacted with the following items: {inters}, please recommend a next item that you think would be suitable for the user."
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+#####——16
+prompt = {}
+prompt["instruction"] = "You have obtained the ordered list of user historical interaction items, which is as follows: {inters}. Using this history as a reference, please select the next item to recommend to the user."
+prompt["response"] = "{item}"
+seqrec_prompt.append(prompt)
+
+all_prompt["seqrec"] = seqrec_prompt
+
+
+
+# ========================================================
+# Task 2 -- Item2Index -- 19 Prompt
+# ========================================================
+# Remove periods when inputting
+
+item2index_prompt = []
+
+# ========================================================
+# Title2Index
+
+#####——0
+prompt = {}
+prompt["instruction"] = "Which item has the title: \"{title}\"?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+#####——1
+prompt = {}
+prompt["instruction"] = "Which item is assigned the title: \"{title}\"?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+#####——2
+prompt = {}
+prompt["instruction"] = "An item is called \"{title}\", could you please let me know which item it is?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+#####——3
+prompt = {}
+prompt["instruction"] = "Which item is called \"{title}\"?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+#####——4
+prompt = {}
+prompt["instruction"] = "One of the items is named \"{title}\", can you tell me which item this is?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+#####——5
+prompt = {}
+prompt["instruction"] = "What is the item that goes by the title \"{title}\"?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+# prompt = {}
+# prompt["instruction"] = "Which item is referred to as \"{title}\"?"
+# prompt["response"] = "{item}"
+# item2index_prompt.append(prompt)
+
+# ========================================================
+# Description2Index
+
+#####——6
+prompt = {}
+prompt["instruction"] = "An item can be described as follows: \"{description}\". Which item is it describing?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+#####——7
+prompt = {}
+prompt["instruction"] = "Can you tell me what item is described as \"{description}\"?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+#####——8
+prompt = {}
+prompt["instruction"] = "Can you provide the item that corresponds to the following description: \"{description}\"?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+
+# prompt = {}
+# prompt["instruction"] = "What is the item described as follows: \"{description}\"?"
+# prompt["response"] = "{item}"
+# item2index_prompt.append(prompt)
+
+#####——9
+prompt = {}
+prompt["instruction"] = "Which item has the following characteristics: \"{description}\"?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+#####——10
+prompt = {}
+prompt["instruction"] = "Which item is characterized by the following description: \"{description}\"?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+#####——11
+prompt = {}
+prompt["instruction"] = "I am curious to know which item can be described as follows: \"{description}\". Can you tell me?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+# ========================================================
+# Title and Description to index
+
+#####——12
+prompt = {}
+prompt["instruction"] = "An item is called \"{title}\" and described as \"{description}\", can you tell me which item it is?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+#####——13
+prompt = {}
+prompt["instruction"] = "Could you please identify what item is called \"{title}\" and described as \"{description}\"?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+#####——14
+prompt = {}
+prompt["instruction"] = "Which item is called \"{title}\" and has the characteristics described below: \"{description}\"?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+#####——15
+prompt = {}
+prompt["instruction"] = "Please show me which item is named \"{title}\" and its corresponding description is: \"{description}\"."
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+
+# prompt = {}
+# prompt["instruction"] = "Here is an item called \"{title}\" and described as \"{description}\". Which item is it?"
+# prompt["response"] = "{item}"
+# item2index_prompt.append(prompt)
+
+#####——16
+prompt = {}
+prompt["instruction"] = "Determine which item this is by its title and description. The title is: \"{title}\", and the description is: \"{description}\"."
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+#####——17
+prompt = {}
+prompt["instruction"] = "Based on the title: \"{title}\", and the description: \"{description}\", answer which item is this?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+#####——18
+prompt = {}
+prompt["instruction"] = "Can you identify the item from the provided title: \"{title}\", and description: \"{description}\"?"
+prompt["response"] = "{item}"
+item2index_prompt.append(prompt)
+
+all_prompt["item2index"] = item2index_prompt
+
+
+# ========================================================
+# Task 3 -- Index2Item --17 Prompt
+# ========================================================
+# Remove periods when inputting
+
+index2item_prompt = []
+
+# ========================================================
+# Index2Title
+
+#####——0
+prompt = {}
+prompt["instruction"] = "What is the title of item {item}?"
+prompt["response"] = "{title}"
+index2item_prompt.append(prompt)
+
+#####——1
+prompt = {}
+prompt["instruction"] = "What title is assigned to item {item}?"
+prompt["response"] = "{title}"
+index2item_prompt.append(prompt)
+
+#####——2
+prompt = {}
+prompt["instruction"] = "Could you please tell me what item {item} is called?"
+prompt["response"] = "{title}"
+index2item_prompt.append(prompt)
+
+#####——3
+prompt = {}
+prompt["instruction"] = "Can you provide the title of item {item}?"
+prompt["response"] = "{title}"
+index2item_prompt.append(prompt)
+
+#####——4
+prompt = {}
+prompt["instruction"] = "What item {item} is referred to as?"
+prompt["response"] = "{title}"
+index2item_prompt.append(prompt)
+
+#####——5
+prompt = {}
+prompt["instruction"] = "Would you mind informing me about the title of item {item}?"
+prompt["response"] = "{title}"
+index2item_prompt.append(prompt)
+
+# ========================================================
+# Index2Description
+
+#####——6
+prompt = {}
+prompt["instruction"] = "Please provide a description of item {item}."
+prompt["response"] = "{description}"
+index2item_prompt.append(prompt)
+
+#####——7
+prompt = {}
+prompt["instruction"] = "Briefly describe item {item}."
+prompt["response"] = "{description}"
+index2item_prompt.append(prompt)
+
+#####——8
+prompt = {}
+prompt["instruction"] = "Can you share with me the description corresponding to item {item}?"
+prompt["response"] = "{description}"
+index2item_prompt.append(prompt)
+
+#####——9
+prompt = {}
+prompt["instruction"] = "What is the description of item {item}?"
+prompt["response"] = "{description}"
+index2item_prompt.append(prompt)
+
+#####——10
+prompt = {}
+prompt["instruction"] = "How to describe the characteristics of item {item}?"
+prompt["response"] = "{description}"
+index2item_prompt.append(prompt)
+
+#####——11
+prompt = {}
+prompt["instruction"] = "Could you please tell me what item {item} looks like?"
+prompt["response"] = "{description}"
+index2item_prompt.append(prompt)
+
+
+# ========================================================
+# index to Title and Description
+
+#####——12
+prompt = {}
+prompt["instruction"] = "What is the title and description of item {item}?"
+prompt["response"] = "{title}\n\n{description}"
+index2item_prompt.append(prompt)
+
+#####——13
+prompt = {}
+prompt["instruction"] = "Can you provide the corresponding title and description for item {item}?"
+prompt["response"] = "{title}\n\n{description}"
+index2item_prompt.append(prompt)
+
+#####——14
+prompt = {}
+prompt["instruction"] = "Please tell me what item {item} is called, along with a brief description of it."
+prompt["response"] = "{title}\n\n{description}"
+index2item_prompt.append(prompt)
+
+#####——15
+prompt = {}
+prompt["instruction"] = "Would you mind informing me about the title of the item {item} and how to describe its characteristics?"
+prompt["response"] = "{title}\n\n{description}"
+index2item_prompt.append(prompt)
+
+#####——16
+prompt = {}
+prompt["instruction"] = "I need to know the title and description of item {item}. Could you help me with that?"
+prompt["response"] = "{title}\n\n{description}"
+index2item_prompt.append(prompt)
+
+all_prompt["index2item"] = index2item_prompt
+
+
+
+
+
+# ========================================================
+# Task 4 -- FusionSequentialRec -- Prompt
+# ========================================================
+
+
+fusionseqrec_prompt = []
+
+#####——0
+prompt = {}
+prompt["instruction"] = "The user has sequentially interacted with items {inters}. Can you recommend the next item for him? Tell me the title of the item？"
+prompt["response"] = "{title}"
+fusionseqrec_prompt.append(prompt)
+
+#####——1
+prompt = {}
+prompt["instruction"] = "Based on the user's historical interactions: {inters}, try to predict the title of the item that the user may need next."
+prompt["response"] = "{title}"
+fusionseqrec_prompt.append(prompt)
+
+#####——2
+prompt = {}
+prompt["instruction"] = "Utilizing the user's past ordered interactions, which include items {inters}, please recommend the next item you think is suitable for the user and provide its title."
+prompt["response"] = "{title}"
+fusionseqrec_prompt.append(prompt)
+
+
+#####——3
+prompt = {}
+prompt["instruction"] = "After interacting with items {inters}, what is the most probable item for the user to interact with next? Kindly provide the item's title."
+prompt["response"] = "{title}"
+fusionseqrec_prompt.append(prompt)
+
+
+
+
+
+#####——4
+prompt = {}
+prompt["instruction"] = "Please review the user's historical interactions: {inters}, and describe what kind of item he still needs."
+prompt["response"] = "{description}"
+fusionseqrec_prompt.append(prompt)
+
+#####——5
+prompt = {}
+prompt["instruction"] = "Here is the item interaction history of the user: {inters}, please tell me what features he expects from his next item."
+prompt["response"] = "{description}"
+fusionseqrec_prompt.append(prompt)
+
+#####——6
+prompt = {}
+prompt["instruction"] = "By analyzing the user's historical interactions with items {inters}, can you infer what the user's next interactive item will look like?"
+prompt["response"] = "{description}"
+fusionseqrec_prompt.append(prompt)
+
+#####——7
+prompt = {}
+prompt["instruction"] = "Access the user's historical item interaction records: {inters}. Your objective is to describe the next potential item for him, taking into account his past interactions."
+prompt["response"] = "{description}"
+fusionseqrec_prompt.append(prompt)
+
+
+
+
+
+
+#####——8
+prompt = {}
+prompt["instruction"] = "Given the title sequence of user historical interactive items: {inter_titles}, can you recommend a suitable next item for the user?"
+prompt["response"] = "{item}"
+fusionseqrec_prompt.append(prompt)
+
+#####——9
+prompt = {}
+prompt["instruction"] = "I possess a user's past interaction history, denoted by the title sequence of interactive items: {inter_titles}, and I am interested in knowing the user's next most desired item. Can you help me?"
+prompt["response"] = "{item}"
+fusionseqrec_prompt.append(prompt)
+
+#####——10
+prompt = {}
+prompt["instruction"] = "Considering the title sequence of user history interaction items: {inter_titles}. What is the next recommendation for the user?"
+prompt["response"] = "{item}"
+fusionseqrec_prompt.append(prompt)
+
+#####——11
+prompt = {}
+prompt["instruction"] = "You have obtained the ordered title list of user historical interaction items, as follows: {inter_titles}. Based on this historical context, kindly choose the subsequent item for user recommendation."
+prompt["response"] = "{item}"
+fusionseqrec_prompt.append(prompt)
+
+
+all_prompt["fusionseqrec"] = fusionseqrec_prompt
+
+
+
+
+
+
+
+# ========================================================
+# Task 5 -- ItemSearch -- Prompt
+# ========================================================
+
+
+itemsearch_prompt = []
+
+#####——0
+prompt = {}
+prompt["instruction"] = "Here is the historical interactions of a user: {inters}. And his personalized preferences are as follows: \"{explicit_preference}\". Your task is to recommend an item that is consistent with the user's preference."
+prompt["response"] = "{item}"
+itemsearch_prompt.append(prompt)
+
+#####——1
+prompt = {}
+prompt["instruction"] = "The user has interacted with a list of items, which are as follows: {inters}. Based on these interacted items, the user current intent is as follows \"{user_related_intention}\", and your task is to generate an item that matches the user's current intent."
+prompt["response"] = "{item}"
+itemsearch_prompt.append(prompt)
+
+#####——2
+prompt = {}
+prompt["instruction"] = "As a recommender system, you are assisting a user who has recently interacted with the following items: {inters}. The user expresses a desire to obtain another item with the following characteristics: \"{item_related_intention}\". Please recommend an item that meets these criteria."
+prompt["response"] = "{item}"
+itemsearch_prompt.append(prompt)
+
+#####——3
+prompt = {}
+prompt["instruction"] = "Using the user's current query: \"{query}\" and his historical interactions: {inters}, you can estimate the user's preferences \"{explicit_preference}\". Please respond to the user's query by selecting an item that best matches his preference and query."
+prompt["response"] = "{item}"
+itemsearch_prompt.append(prompt)
+
+#####——4
+prompt = {}
+prompt["instruction"] = "The user needs a new item and searches for: \"{query}\". In addition, he has previously interacted with: {inters}. You can obtain his preference by analyzing his historical interactions: \"{explicit_preference}\". Can you recommend an item that best matches the search query and preferences?"
+prompt["response"] = "{item}"
+itemsearch_prompt.append(prompt)
+
+#####——5
+prompt = {}
+prompt["instruction"] = "Based on the user's historical interactions with the following items: {inters}. You can infer his preference by observing the historical interactions: \"{explicit_preference}\". Now the user wants a new item and searches for: \"{query}\". Please select a suitable item that matches his preference and search intent."
+prompt["response"] = "{item}"
+itemsearch_prompt.append(prompt)
+
+
+
+
+
+#####——6
+prompt = {}
+prompt["instruction"] = "Suppose you are a search engine, now a user searches that: \"{query}\", can you select an item to respond to the user's query?"
+prompt["response"] = "{item}"
+itemsearch_prompt.append(prompt)
+
+#####——7
+prompt = {}
+prompt["instruction"] = "As a search engine, your task is to answer the user's query by generating a related item. The user's query is provided as \"{query}\". Please provide your generated item as your answer."
+prompt["response"] = "{item}"
+itemsearch_prompt.append(prompt)
+
+#####——8
+prompt = {}
+prompt["instruction"] = "As a recommender system, your task is to recommend an item that is related to the user's request, which is specified as follows: \"{query}\". Please provide your recommendation."
+prompt["response"] = "{item}"
+itemsearch_prompt.append(prompt)
+
+#####——9
+prompt = {}
+prompt["instruction"] = "You meet a user's query: \"{query}\". Please respond to this user by selecting an appropriate item."
+prompt["response"] = "{item}"
+itemsearch_prompt.append(prompt)
+
+
+#####——10
+prompt = {}
+prompt["instruction"] = "Your task is to recommend the best item that matches the user's query. Here is the search query of the user: \"{query}\", tell me the item you recommend."
+prompt["response"] = "{item}"
+itemsearch_prompt.append(prompt)
+
+all_prompt["itemsearch"] = itemsearch_prompt
+
+
+
+
+
+# ========================================================
+# Task 6 -- PreferenceObtain -- Prompt
+# ========================================================
+
+preferenceobtain_prompt = []
+
+#####——0
+prompt = {}
+prompt["instruction"] = "The user has interacted with items {inters} in chronological order. Please estimate his preferences."
+prompt["response"] = "{explicit_preference}"
+preferenceobtain_prompt.append(prompt)
+
+#####——1
+prompt = {}
+prompt["instruction"] = "Based on the items that the user has interacted with: {inters}, can you infer what preferences he has?"
+prompt["response"] = "{explicit_preference}"
+preferenceobtain_prompt.append(prompt)
+
+#####——3
+prompt = {}
+prompt["instruction"] = "Can you provide a summary of the user's preferences based on his historical interactions: {inters}?"
+prompt["response"] = "{explicit_preference}"
+preferenceobtain_prompt.append(prompt)
+
+#####——4
+prompt = {}
+prompt["instruction"] = "After interacting with items {inters} in order, what preferences do you think the user has?"
+prompt["response"] = "{explicit_preference}"
+preferenceobtain_prompt.append(prompt)
+
+#####——5
+prompt = {}
+prompt["instruction"] = "Here is the item interaction history of the user: {inters}, could you please infer the user's preferences."
+prompt["response"] = "{explicit_preference}"
+preferenceobtain_prompt.append(prompt)
+
+#####——6
+prompt = {}
+prompt["instruction"] = "Based on the user's historical interaction records: {inters}, what are your speculations about his preferences?"
+prompt["response"] = "{explicit_preference}"
+preferenceobtain_prompt.append(prompt)
+
+#####——7
+prompt = {}
+prompt["instruction"] = "Given the user's historical interactive items arranged in chronological order: {inters}, what can be inferred about the preferences of the user?"
+prompt["response"] = "{explicit_preference}"
+preferenceobtain_prompt.append(prompt)
+
+#####——8
+prompt = {}
+prompt["instruction"] = "Can you speculate on the user's preferences based on his historical item interaction records: {inters}?"
+prompt["response"] = "{explicit_preference}"
+preferenceobtain_prompt.append(prompt)
+
+#####——9
+prompt = {}
+prompt["instruction"] = "What is the preferences of a user who has previously interacted with items {inters} sequentially?"
+prompt["response"] = "{explicit_preference}"
+preferenceobtain_prompt.append(prompt)
+
+#####——10
+prompt = {}
+prompt["instruction"] = "Using the user's historical interactions as input data, summarize the user's preferences. The historical interactions are provided as follows: {inters}."
+prompt["response"] = "{explicit_preference}"
+preferenceobtain_prompt.append(prompt)
+
+#####——11
+prompt = {}
+prompt["instruction"] = "Utilizing the ordered list of the user's historical interaction items as a reference, please make an informed estimation of the user's preferences. The historical interactions are as follows: {inters}."
+prompt["response"] = "{explicit_preference}"
+preferenceobtain_prompt.append(prompt)
+
+all_prompt["preferenceobtain"] = preferenceobtain_prompt

run.sh

@@ -0,0 +1,98 @@
+export WANDB_MODE=disabled
+export CUDA_LAUNCH_BLOCKING=1
+
+DATASET=Games
+BASE_MODEL=huggyllama/llama-7b
+DATA_PATH=./data
+OUTPUT_DIR=./ckpt/$DATASET/
+
+torchrun --nproc_per_node=8 --master_port=3324 finetune.py \
+    --base_model $BASE_MODEL \
+    --output_dir $OUTPUT_DIR \
+    --dataset $DATASET \
+    --data_path $DATA_PATH \
+    --per_device_batch_size 8 \
+    --gradient_accumulation_steps 2 \
+    --learning_rate 5e-5 \
+    --epochs 4 \
+    --weight_decay 0.01 \
+    --save_and_eval_strategy epoch \
+    --deepspeed ./config/ds_z3_bf16.json \
+    --bf16 \
+    --only_train_response \
+    --tasks seqrec,item2index,index2item,fusionseqrec,itemsearch,preferenceobtain \
+    --train_prompt_sample_num 1,1,1,1,1,1 \
+    --train_data_sample_num 0,0,0,100000,0,0 \
+    --index_file .index.json
+
+
+cd convert
+nohup ./convert.sh $OUTPUT_DIR >convert.log 2>&1 &
+cd ..
+
+
+
+
+
+
+DATASET=Arts
+BASE_MODEL=huggyllama/llama-7b
+DATA_PATH=./data
+OUTPUT_DIR=./ckpt/$DATASET/
+
+torchrun --nproc_per_node=8 --master_port=3324 finetune.py \
+    --base_model $BASE_MODEL \
+    --output_dir $OUTPUT_DIR \
+    --dataset $DATASET \
+    --data_path $DATA_PATH \
+    --per_device_batch_size 8 \
+    --gradient_accumulation_steps 2 \
+    --learning_rate 5e-5 \
+    --epochs 4 \
+    --weight_decay 0.01 \
+    --save_and_eval_strategy epoch \
+    --deepspeed ./config/ds_z3_bf16.json \
+    --bf16 \
+    --only_train_response \
+    --tasks seqrec,item2index,index2item,fusionseqrec,itemsearch,preferenceobtain \
+    --train_prompt_sample_num 1,1,1,1,1,1 \
+    --train_data_sample_num 0,0,0,30000,0,0 \
+    --index_file .index.json
+
+
+cd convert
+nohup ./convert.sh $OUTPUT_DIR >convert.log 2>&1 &
+cd ..
+
+
+
+
+
+DATASET=Instruments
+BASE_MODEL=huggyllama/llama-7b
+DATA_PATH=./data
+OUTPUT_DIR=./ckpt/$DATASET/
+
+torchrun --nproc_per_node=8 --master_port=3324 finetune.py \
+    --base_model $BASE_MODEL \
+    --output_dir $OUTPUT_DIR \
+    --dataset $DATASET \
+    --data_path $DATA_PATH \
+    --per_device_batch_size 8 \
+    --gradient_accumulation_steps 2 \
+    --learning_rate 5e-5 \
+    --epochs 4 \
+    --weight_decay 0.01 \
+    --save_and_eval_strategy epoch \
+    --deepspeed ./config/ds_z3_bf16.json \
+    --bf16 \
+    --only_train_response \
+    --tasks seqrec,item2index,index2item,fusionseqrec,itemsearch,preferenceobtain \
+    --train_prompt_sample_num 1,1,1,1,1,1 \
+    --train_data_sample_num 0,0,0,20000,0,0 \
+    --index_file .index.json
+
+
+cd convert
+nohup ./convert.sh $OUTPUT_DIR >convert.log 2>&1 &
+cd ..

run_test.sh

@@ -0,0 +1,17 @@
+
+
+DATASET=Games
+DATA_PATH=./data
+OUTPUT_DIR=./ckpt/$DATASET/
+RESULTS_FILE=./results/$DATASET/xxx.json
+
+python test.py \
+    --gpu_id 0 \
+    --ckpt_path $CKPT_PATH \
+    --dataset $DATASET \
+    --data_path $DATA_PATH \
+    --results_file $RESULTS_FILE \
+    --test_batch_size 1 \
+    --num_beams 20 \
+    --test_prompt_ids all \
+    --index_file .index.json

test.py

@@ -0,0 +1,175 @@
+import argparse
+import json
+import os
+import sys
+from typing import List
+
+import torch
+import transformers
+from peft import PeftModel
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import LlamaForCausalLM, LlamaTokenizer, LlamaConfig
+
+from utils import *
+from collator import TestCollator
+from prompt import all_prompt
+from evaluate import get_topk_results, get_metrics_results
+
+
+def test(args):
+
+    set_seed(args.seed)
+    print(vars(args))
+
+    device_map = {"": args.gpu_id}
+    device = torch.device("cuda",args.gpu_id)
+
+
+    tokenizer = LlamaTokenizer.from_pretrained(args.ckpt_path)
+    if args.lora:
+        model = LlamaForCausalLM.from_pretrained(
+            args.base_model,
+            torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+            device_map=device_map,
+        )
+        model.resize_token_embeddings(len(tokenizer))
+        model = PeftModel.from_pretrained(
+            model,
+            args.ckpt_path,
+            torch_dtype=torch.bfloat16,
+            device_map=device_map,
+        )
+    else:
+        model = LlamaForCausalLM.from_pretrained(
+            args.ckpt_path,
+            torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True,
+            device_map=device_map,
+        )
+    # assert model.config.vocab_size == len(tokenizer)
+
+    if args.test_prompt_ids == "all":
+        if args.test_task.lower() == "seqrec":
+            prompt_ids = range(len(all_prompt["seqrec"]))
+        elif args.test_task.lower() == "itemsearch":
+            prompt_ids = range(len(all_prompt["itemsearch"]))
+        elif args.test_task.lower() == "fusionseqrec":
+            prompt_ids = range(len(all_prompt["fusionseqrec"]))
+    else:
+        prompt_ids = [int(_) for _ in args.test_prompt_ids.split(",")]
+
+    test_data = load_test_dataset(args)
+    collator = TestCollator(args, tokenizer)
+    all_items = test_data.get_all_items()
+
+
+    prefix_allowed_tokens = test_data.get_prefix_allowed_tokens_fn(tokenizer)
+
+    test_loader = DataLoader(test_data, batch_size=args.test_batch_size, collate_fn=collator,
+                             shuffle=True, num_workers=4, pin_memory=True)
+
+
+    print("data num:", len(test_data))
+
+    model.eval()
+
+    metrics = args.metrics.split(",")
+    all_prompt_results = []
+    with torch.no_grad():
+        for prompt_id in prompt_ids:
+
+            test_loader.dataset.set_prompt(prompt_id)
+            metrics_results = {}
+            total = 0
+
+            for step, batch in enumerate(tqdm(test_loader)):
+                inputs = batch[0].to(device)
+                targets = batch[1]
+                total += len(targets)
+
+                output = model.generate(
+                    input_ids=inputs["input_ids"],
+                    attention_mask=inputs["attention_mask"],
+                    max_new_tokens=10,
+                    # max_length=10,
+                    prefix_allowed_tokens_fn=prefix_allowed_tokens,
+                    num_beams=args.num_beams,
+                    num_return_sequences=args.num_beams,
+                    output_scores=True,
+                    return_dict_in_generate=True,
+                    early_stopping=True,
+                )
+                output_ids = output["sequences"]
+                scores = output["sequences_scores"]
+
+                output = tokenizer.batch_decode(
+                    output_ids, skip_special_tokens=True
+                )
+                # print(output)
+                topk_res = get_topk_results(output,scores,targets,args.num_beams,
+                                            all_items=all_items if args.filter_items else None)
+
+                batch_metrics_res = get_metrics_results(topk_res, metrics)
+                # print(batch_metrics_res)
+
+                for m, res in batch_metrics_res.items():
+                    if m not in metrics_results:
+                        metrics_results[m] = res
+                    else:
+                        metrics_results[m] += res
+
+                if (step+1)%10 == 0:
+                    temp={}
+                    for m in metrics_results:
+                        temp[m] = metrics_results[m] / total
+                    print(temp)
+
+            for m in metrics_results:
+                metrics_results[m] = metrics_results[m] / total
+
+            all_prompt_results.append(metrics_results)
+            print("======================================================")
+            print("Prompt {} results: ".format(prompt_id), metrics_results)
+            print("======================================================")
+            print("")
+
+    mean_results = {}
+    min_results = {}
+    max_results = {}
+
+    for m in metrics:
+        all_res = [_[m] for _ in all_prompt_results]
+        mean_results[m] = sum(all_res)/len(all_res)
+        min_results[m] = min(all_res)
+        max_results[m] = max(all_res)
+
+    print("======================================================")
+    print("Mean results: ", mean_results)
+    print("Min results: ", min_results)
+    print("Max results: ", max_results)
+    print("======================================================")
+
+
+    save_data={}
+    save_data["test_prompt_ids"] = args.test_prompt_ids
+    save_data["mean_results"] = mean_results
+    save_data["min_results"] = min_results
+    save_data["max_results"] = max_results
+    save_data["all_prompt_results"] = all_prompt_results
+
+    with open(args.results_file, "w") as f:
+        json.dump(save_data, f, indent=4)
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="LLMRec_test")
+    parser = parse_global_args(parser)
+    parser = parse_dataset_args(parser)
+    parser = parse_test_args(parser)
+
+    args = parser.parse_args()
+
+    test(args)

test_ddp.py

@@ -0,0 +1,238 @@
+import argparse
+import json
+import os
+import sys
+
+import torch
+import transformers
+import torch.distributed as dist
+from torch.utils.data.distributed import DistributedSampler
+from torch.nn.parallel import DistributedDataParallel
+from peft import PeftModel
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+from transformers import LlamaForCausalLM, LlamaTokenizer, LlamaConfig
+
+from utils import *
+from collator import TestCollator
+from prompt import all_prompt
+from evaluate import get_topk_results, get_metrics_results
+
+
+def test_ddp(args):
+
+    set_seed(args.seed)
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    local_rank = int(os.environ.get("LOCAL_RANK") or 0)
+    torch.cuda.set_device(local_rank)
+    if local_rank == 0:
+        print(vars(args))
+
+    dist.init_process_group(backend="nccl", world_size=world_size, rank=local_rank)
+
+    device_map = {"": local_rank}
+    device = torch.device("cuda",local_rank)
+
+    tokenizer = LlamaTokenizer.from_pretrained(args.ckpt_path)
+    args.lora=True
+    if args.lora:
+        model = LlamaForCausalLM.from_pretrained(
+            args.base_model,
+            torch_dtype=torch.float16,
+            low_cpu_mem_usage=True,
+            device_map=device_map,
+        )
+        model.resize_token_embeddings(len(tokenizer))
+        model = PeftModel.from_pretrained(
+            model,
+            args.ckpt_path,
+            torch_dtype=torch.float16,
+            device_map=device_map,
+        )
+    else:
+        model = LlamaForCausalLM.from_pretrained(
+            args.ckpt_path,
+            torch_dtype=torch.float16,
+            low_cpu_mem_usage=True,
+            device_map=device_map,
+        )
+    # assert model.config.vocab_size == len(tokenizer)
+    model = DistributedDataParallel(model, device_ids=[local_rank])
+
+    if args.test_prompt_ids == "all":
+        if args.test_task.lower() == "seqrec":
+            prompt_ids = range(len(all_prompt["seqrec"]))
+        elif args.test_task.lower() == "itemsearch":
+            prompt_ids = range(len(all_prompt["itemsearch"]))
+        elif args.test_task.lower() == "fusionseqrec":
+            prompt_ids = range(len(all_prompt["fusionseqrec"]))
+    else:
+        prompt_ids = [int(_) for _ in args.test_prompt_ids.split(",")]
+
+    test_data = load_test_dataset(args)
+    ddp_sampler = DistributedSampler(test_data, num_replicas=world_size, rank=local_rank, drop_last=True)
+
+    test_data = load_test_dataset(args)
+    collator = TestCollator(args, tokenizer)
+    all_items = test_data.get_all_items()
+
+
+    prefix_allowed_tokens = test_data.get_prefix_allowed_tokens_fn(tokenizer)
+
+
+    test_loader = DataLoader(test_data, batch_size=args.test_batch_size, collate_fn=collator,
+                             sampler=ddp_sampler, num_workers=2, pin_memory=True)
+
+    if local_rank == 0:
+        print("data num:", len(test_data))
+
+    model.eval()
+
+    metrics = args.metrics.split(",")
+    all_prompt_results = []
+    with torch.no_grad():
+
+        for prompt_id in prompt_ids:
+
+            if local_rank == 0:
+                print("Start prompt: ",prompt_id)
+
+            test_loader.dataset.set_prompt(prompt_id)
+            metrics_results = {}
+            total = 0
+
+            for step, batch in enumerate(tqdm(test_loader)):
+                inputs = batch[0].to(device)
+                targets = batch[1]
+                bs = len(targets)
+                num_beams = args.num_beams
+                while True:
+                    try:
+                        output = model.module.generate(
+                            input_ids=inputs["input_ids"],
+                            attention_mask=inputs["attention_mask"],
+                            max_new_tokens=10,
+                            prefix_allowed_tokens_fn=prefix_allowed_tokens,
+                            num_beams=num_beams,
+                            num_return_sequences=num_beams,
+                            output_scores=True,
+                            return_dict_in_generate=True,
+                            early_stopping=True,
+                        )
+                        break
+                    except torch.cuda.OutOfMemoryError as e:
+                        print("Out of memory!")
+                        num_beams = num_beams -1
+                        print("Beam:", num_beams)
+                    except Exception:
+                        raise RuntimeError
+
+                output_ids = output["sequences"]
+                scores = output["sequences_scores"]
+
+                output = tokenizer.batch_decode(
+                    output_ids, skip_special_tokens=True
+                )
+
+                topk_res = get_topk_results(output, scores, targets, num_beams,
+                                            all_items=all_items if args.filter_items else None)
+
+                bs_gather_list = [None for _ in range(world_size)]
+                dist.all_gather_object(obj=bs, object_list=bs_gather_list)
+                total += sum(bs_gather_list)
+                res_gather_list = [None for _ in range(world_size)]
+                dist.all_gather_object(obj=topk_res, object_list=res_gather_list)
+
+
+                if local_rank == 0:
+                    all_device_topk_res = []
+                    for ga_res in res_gather_list:
+                        all_device_topk_res += ga_res
+                    batch_metrics_res = get_metrics_results(all_device_topk_res, metrics)
+                    for m, res in batch_metrics_res.items():
+                        if m not in metrics_results:
+                            metrics_results[m] = res
+                        else:
+                            metrics_results[m] += res
+
+                    if (step + 1) % 50 == 0:
+                        temp = {}
+                        for m in metrics_results:
+                            temp[m] = metrics_results[m] / total
+                        print(temp)
+
+                dist.barrier()
+
+            if local_rank == 0:
+                for m in metrics_results:
+                    metrics_results[m] = metrics_results[m] / total
+
+                all_prompt_results.append(metrics_results)
+                print("======================================================")
+                print("Prompt {} results: ".format(prompt_id), metrics_results)
+                print("======================================================")
+                print("")
+
+            dist.barrier()
+
+    dist.barrier()
+
+    if local_rank == 0:
+        mean_results = {}
+        min_results = {}
+        max_results = {}
+
+        for m in metrics:
+            all_res = [_[m] for _ in all_prompt_results]
+            mean_results[m] = sum(all_res)/len(all_res)
+            min_results[m] = min(all_res)
+            max_results[m] = max(all_res)
+
+        print("======================================================")
+        print("Mean results: ", mean_results)
+        print("Min results: ", min_results)
+        print("Max results: ", max_results)
+        print("======================================================")
+
+
+        save_data={}
+        save_data["test_prompt_ids"] = args.test_prompt_ids
+        save_data["mean_results"] = mean_results
+        save_data["min_results"] = min_results
+        save_data["max_results"] = max_results
+        save_data["all_prompt_results"] = all_prompt_results
+
+        with open(args.results_file, "w") as f:
+            json.dump(save_data, f, indent=4)
+        print("Save file: ", args.results_file)
+
+        import smtplib
+        from email.mime.text import MIMEText
+        mail_host = 'smtp.qq.com'
+        mail_code = 'ouzplpngooqndjcb'
+        sender = '[email protected]'
+        receiver = '[email protected]'
+
+        task = '[v67: evaluate lcrec]'
+        message = MIMEText('Task {task} Finished'.format(task = task), 'plain', 'utf-8')
+        message['Subject'] = 'Auto Email'
+        message['From'] = sender
+        message['To'] = receiver
+
+        server = smtplib.SMTP_SSL("smtp.qq.com", 465)
+        server.login(sender, mail_code)
+        server.sendmail(sender, receiver, message.as_string())
+
+        server.quit()
+
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="LLMRec_test")
+    parser = parse_global_args(parser)
+    parser = parse_dataset_args(parser)
+    parser = parse_test_args(parser)
+
+    args = parser.parse_args()
+
+    test_ddp(args)

test_ddp.sh

@@ -0,0 +1,14 @@
+DATASET=Instruments
+DATA_PATH=$datain/v-yinju/rqvae-zzx/data
+CKPT_PATH=$datain/v-yinju/rq-llama
+RESULTS_FILE=$CKPT_PATH/result.json
+
+torchrun --nproc_per_node=8 --master_port=4324 test_ddp.py \
+    --ckpt_path $CKPT_PATH \
+    --dataset $DATASET \
+    --data_path $DATA_PATH \
+    --results_file $RESULTS_FILE \
+    --test_batch_size 1 \
+    --num_beams 20 \
+    --test_prompt_ids all \
+    --index_file .index.json

utils.py

@@ -0,0 +1,196 @@
+import json
+import logging
+import os
+import random
+import datetime
+
+import numpy as np
+import torch
+from torch.utils.data import ConcatDataset
+from data import SeqRecDataset, ItemFeatDataset, ItemSearchDataset, FusionSeqRecDataset, SeqRecTestDataset, PreferenceObtainDataset
+
+
+def parse_global_args(parser):
+
+
+    parser.add_argument("--seed", type=int, default=42, help="Random seed")
+
+    parser.add_argument("--base_model", type=str,
+                        default="./llama-7b/",
+                        help="basic model path")
+    parser.add_argument("--output_dir", type=str,
+                        default="./ckpt/",
+                        help="The output directory")
+
+
+    return parser
+
+def parse_dataset_args(parser):
+    parser.add_argument("--data_path", type=str, default="",
+                        help="data directory")
+    parser.add_argument("--tasks", type=str, default="seqrec,item2index,index2item,fusionseqrec,itemsearch,preferenceobtain",
+                        help="Downstream tasks, separate by comma")
+    parser.add_argument("--dataset", type=str, default="Games", help="Dataset name")
+    parser.add_argument("--index_file", type=str, default=".index.json", help="the item indices file")
+    parser.add_argument("--dataloader_num_workers", type=int, default=0, help="dataloader num_workers")
+    parser.add_argument("--dataloader_prefetch_factor", type=int, default=2, help="dataloader prefetch_factor")
+    
+
+    # arguments related to sequential task
+    parser.add_argument("--max_his_len", type=int, default=20,
+                        help="the max number of items in history sequence, -1 means no limit")
+    parser.add_argument("--add_prefix", action="store_true", default=False,
+                        help="whether add sequential prefix in history")
+    parser.add_argument("--his_sep", type=str, default=", ", help="The separator used for history")
+    parser.add_argument("--only_train_response", action="store_true", default=False,
+                        help="whether only train on responses")
+
+    parser.add_argument("--train_prompt_sample_num", type=str, default="1,1,1,1,1,1",
+                        help="the number of sampling prompts for each task")
+    parser.add_argument("--train_data_sample_num", type=str, default="0,0,0,100000,0,0",
+                        help="the number of sampling prompts for each task")
+
+    parser.add_argument("--valid_prompt_id", type=int, default=0,
+                        help="The prompt used for validation")
+    parser.add_argument("--sample_valid", action="store_true", default=True,
+                        help="use sampled prompt for validation")
+    parser.add_argument("--valid_prompt_sample_num", type=int, default=2,
+                        help="the number of sampling validation sequential recommendation prompts")
+
+    return parser
+
+def parse_train_args(parser):
+
+    parser.add_argument("--optim", type=str, default="adamw_torch", help='The name of the optimizer')
+    parser.add_argument("--epochs", type=int, default=4)
+    parser.add_argument("--learning_rate", type=float, default=2e-5)
+    parser.add_argument("--per_device_batch_size", type=int, default=8)
+    parser.add_argument("--gradient_accumulation_steps", type=int, default=2)
+    parser.add_argument("--logging_step", type=int, default=10)
+    parser.add_argument("--model_max_length", type=int, default=2048)
+    parser.add_argument("--weight_decay", type=float, default=0.01)
+
+    parser.add_argument("--lora_r", type=int, default=8)
+    parser.add_argument("--lora_alpha", type=int, default=32)
+    parser.add_argument("--lora_dropout", type=float, default=0.05)
+    parser.add_argument("--lora_target_modules", type=str,
+                        default="q_proj,v_proj,k_proj,o_proj,gate_proj,down_proj,up_proj", help="separate by comma")
+    parser.add_argument("--lora_modules_to_save", type=str,
+                        default="embed_tokens,lm_head", help="separate by comma")
+
+    parser.add_argument("--resume_from_checkpoint", type=str, default=None, help="either training checkpoint or final adapter")
+
+    parser.add_argument("--warmup_ratio", type=float, default=0.01)
+    parser.add_argument("--lr_scheduler_type", type=str, default="cosine")
+    parser.add_argument("--save_and_eval_strategy", type=str, default="epoch")
+    parser.add_argument("--save_and_eval_steps", type=int, default=1000)
+    parser.add_argument("--fp16",  action="store_true", default=False)
+    parser.add_argument("--bf16", action="store_true", default=False)
+    parser.add_argument("--deepspeed", type=str, default="./config/ds_z3_bf16.json")
+
+    return parser
+
+def parse_test_args(parser):
+
+    parser.add_argument("--ckpt_path", type=str,
+                        default="",
+                        help="The checkpoint path")
+    parser.add_argument("--lora", action="store_true", default=False)
+    parser.add_argument("--filter_items", action="store_true", default=False,
+                        help="whether filter illegal items")
+
+    parser.add_argument("--results_file", type=str,
+                        default="./results/test-ddp.json",
+                        help="result output path")
+
+    parser.add_argument("--test_batch_size", type=int, default=1)
+    parser.add_argument("--num_beams", type=int, default=20)
+    parser.add_argument("--sample_num", type=int, default=-1,
+                        help="test sample number, -1 represents using all test data")
+    parser.add_argument("--gpu_id", type=int, default=0,
+                        help="GPU ID when testing with single GPU")
+    parser.add_argument("--test_prompt_ids", type=str, default="0",
+                        help="test prompt ids, separate by comma. 'all' represents using all")
+    parser.add_argument("--metrics", type=str, default="hit@1,hit@5,hit@10,ndcg@5,ndcg@10",
+                        help="test metrics, separate by comma")
+    parser.add_argument("--test_task", type=str, default="SeqRec")
+
+
+    return parser
+
+
+def get_local_time():
+    cur = datetime.datetime.now()
+    cur = cur.strftime("%b-%d-%Y_%H-%M-%S")
+
+    return cur
+
+
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.enabled = False
+
+def ensure_dir(dir_path):
+
+    os.makedirs(dir_path, exist_ok=True)
+
+
+def load_datasets(args):
+
+    tasks = args.tasks.split(",")
+
+    train_prompt_sample_num = [int(_) for _ in args.train_prompt_sample_num.split(",")]
+    assert len(tasks) == len(train_prompt_sample_num), "prompt sample number does not match task number"
+    train_data_sample_num = [int(_) for _ in args.train_data_sample_num.split(",")]
+    assert len(tasks) == len(train_data_sample_num), "data sample number does not match task number"
+
+    train_datasets = []
+    for task, prompt_sample_num,data_sample_num in zip(tasks,train_prompt_sample_num,train_data_sample_num):
+        if task.lower() == "seqrec":
+            dataset = SeqRecDataset(args, mode="train", prompt_sample_num=prompt_sample_num, sample_num=data_sample_num)
+
+        elif task.lower() == "item2index" or task.lower() == "index2item":
+            dataset = ItemFeatDataset(args, task=task.lower(), prompt_sample_num=prompt_sample_num, sample_num=data_sample_num)
+
+        elif task.lower() == "fusionseqrec":
+            dataset = FusionSeqRecDataset(args, mode="train", prompt_sample_num=prompt_sample_num, sample_num=data_sample_num)
+
+        elif task.lower() == "itemsearch":
+            dataset = ItemSearchDataset(args, mode="train", prompt_sample_num=prompt_sample_num, sample_num=data_sample_num)
+
+        elif task.lower() == "preferenceobtain":
+            dataset = PreferenceObtainDataset(args, prompt_sample_num=prompt_sample_num, sample_num=data_sample_num)
+
+        else:
+            raise NotImplementedError
+        train_datasets.append(dataset)
+
+    train_data = ConcatDataset(train_datasets)
+
+    valid_data = SeqRecDataset(args,"valid",args.valid_prompt_sample_num)
+
+    return train_data, valid_data
+
+def load_test_dataset(args):
+
+    if args.test_task.lower() == "seqrec":
+        test_data = SeqRecDataset(args, mode="test", sample_num=args.sample_num)
+        # test_data = SeqRecTestDataset(args, sample_num=args.sample_num)
+    elif args.test_task.lower() == "itemsearch":
+        test_data = ItemSearchDataset(args, mode="test", sample_num=args.sample_num)
+    elif args.test_task.lower() == "fusionseqrec":
+        test_data = FusionSeqRecDataset(args, mode="test", sample_num=args.sample_num)
+    else:
+        raise NotImplementedError
+
+    return test_data
+
+def load_json(file):
+    with open(file, 'r') as f:
+        data = json.load(f)
+    return data