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,273 @@
+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 VanillaCollator(object):
+    def __init__(self, args, tokenizer):
+        self.args = args
+        self.tokenizer = tokenizer
+    def __call__(self, data):
+        # print('collator data:',data)
+        '''
+        [{
+        'input_ids': 
+            "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n
+            ### Instruction:\n
+                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.\n\n
+            ### Response:", 
+        'labels': 
+            "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n
+            ### Instruction:\n
+                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.\n\n
+            ### Response:
+                Dunlop guitar picks are a top choice of today's pro musician! Dunlop's wide variety of gauges, shapes, sizes and materials 
+                allows the player to select the exact pick for his/her own particular style of playing. From classic country to nu-metal, 
+                every great player knows that their pick is an integral part of their tone, and Dunlop guitar picks are the picks that more 
+                pros rely on in the studio or on stage. Picks are a grossly underrated accessory. Don't sacrifice your tone...pick Dunlop guitar picks!.",
+        'inters': '341,2804,3895,3893,7064', 
+        'item': 'placeholder', 
+        'task': 'inters2description'
+        }, 
+        {
+        'input_ids': 
+            'Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n
+            ### Instruction:\n
+                Based on the user\'s historical interactions with the following items: {inters}. 
+                You can infer his preference by observing the historical interactions: "The user\'s short-term preferences have shift to heavier picks, 
+                suggesting that He is looking for a heavier sound.". Now the user wants a new item and searches for: "I like the durability and 
+                effectiveness of the picks.". Please select a suitable item that matches his preference and search intent.\n\n
+            ### Response:', 
+        'labels': 
+            'Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n
+            ### Instruction:\n
+                Based on the user\'s historical interactions with the following items: {inters}. 
+                You can infer his preference by observing the historical interactions: "The user\'s short-term preferences have shift to heavier picks, 
+                suggesting that He is looking for a heavier sound.". Now the user wants a new item and searches for: "I like the durability and 
+                effectiveness of the picks.". Please select a suitable item that matches his preference and search intent.\n\n
+            ### Response:{item}', 
+        'inters': '122,469,8918', 
+        'item': '7140', 
+        'task': 'itemsearch'
+        }]
+        '''
+        dict_data = {
+            'input_ids': [],
+            'labels': [],
+            'inters': [],
+            'item': [],
+            'users': [],
+            'user': [],
+            'task': []
+        }
+
+        for d in data:
+            for k in dict_data.keys():
+                if k == 'labels':
+                    dict_data[k].append(d[k] + self.tokenizer.eos_token)
+                else:
+                    dict_data[k].append(d[k])
+
+        return dict_data
+
+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):
+            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)
+
+# RuntimeError: Cannot re-initialize CUDA in forked subprocess. 
+# To use CUDA with multiprocessing, you must use the 'spawn' start method.
+# class ValidCollator(object):
+#     def __init__(self, args, model):
+#         self.args = args
+#         self.model = model
+#         self.only_train_response = args.only_train_response
+#         self.tokenizer = model.tokenizer
+#     def __call__(self, data):
+#         llama_model = self.model.model.get_decoder()
+#         for d in data:
+#             inter_emb_list = []
+#             inter_item_list = d['inters'].split(',')
+#             for inter_item in inter_item_list:
+#                 inter_feature = self.model.item_texts[inter_item]['title'] + ' ' + self.model.item_texts[inter_item]['description']
+#                 inter_id = self.tokenizer(inter_feature, return_tensors = 'pt', padding=True, truncation=True).to(self.model.device)
+#                 inter_emb = llama_model(input_ids = inter_id.input_ids, attention_mask = inter_id.attention_mask)
+#                 inter_emb = inter_emb.last_hidden_state * inter_id.attention_mask.unsqueeze(-1)
+#                 inter_emb = inter_emb.sum(dim=1) / inter_id.attention_mask.sum(dim = -1, keepdim = True)
+#                 inter_emb_list.append(inter_emb.detach())
+#             inter_embs = torch.cat(inter_emb_list, dim = 0)
+#             item_feature = self.model.item_texts[d['item']]['title'] + ' ' + self.model.item_texts[d['item']]['description']
+#             item_ids = self.tokenizer(item_feature, return_tensors = 'pt', padding=True, truncation=True).to(self.model.device)
+#             item_emb = llama_model(input_ids = item_ids.input_ids, attention_mask = item_ids.attention_mask)
+#             item_emb = item_emb.last_hidden_state * item_ids.attention_mask.unsqueeze(-1)
+#             item_emb = item_emb.sum(dim=1) / item_ids.attention_mask.sum(dim = -1, keepdim = True)
+#             item_emb = item_emb.detach()
+
+#             rqids = self.model.rqvae.get_indices(torch.cat([inter_embs, item_emb], dim = 0))
+
+#             inters_rqids = rqids.view(-1, rqids.shape[-1]).cpu().numpy().tolist()[:-1]
+#             item_rqid = rqids.view(-1, rqids.shape[-1]).cpu().numpy().tolist()[-1]
+
+#             text_rqids = {}
+#             code = ''
+#             for rqid in inters_rqids:
+#                 for k, idx in enumerate(rqid):
+#                     code = code + self.model.prefix[k].format(idx)
+#                 code = code + ', '
+#             text_rqids['inters'] = code[:-2]
+#             code = ''
+#             for k, idx in enumerate(item_rqid):
+#                 code = code + self.model.prefix[k].format(idx)
+#             text_rqids['item'] = code
+
+#             d['input_ids'] = d['input_ids'].format(inters = text_rqids['inters'])
+#             d['labels'] = d['labels'].format(inters = text_rqids['inters'], item = text_rqids['item'])
+
+#         input_texts = [d["input_ids"] for d in data]
+#         full_texts = [d["labels"] + self.tokenizer.eos_token for d in data]
+
+#         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:
+#             labels[labels == self.tokenizer.pad_token_id] = -100
+#             labels[torch.where(inputs["labels"] != self.tokenizer.pad_token_id)] = -100
+#         inputs["labels"] = labels
+
+#         return inputs
+
+# RuntimeError: Cannot re-initialize CUDA in forked subprocess. 
+# To use CUDA with multiprocessing, you must use the 'spawn' start method.
+# class TestCollator(object):
+#     def __init__(self, args, model):
+#         self.args = args
+#         self.model = model
+#         self.tokenizer = model.tokenizer
+#         if self.tokenizer.pad_token_id is None:
+#             self.tokenizer.pad_token_id = 0
+#         if isinstance(self.tokenizer, LlamaTokenizer):
+#             self.tokenizer.padding_side = "left"
+    
+#     def __call__(self, data):
+#         llama_model = self.model.model.get_decoder()
+#         for d in data:
+#             inter_emb_list = []
+#             inter_item_list = d['inters'].split(',')
+#             for inter_item in inter_item_list:
+#                 inter_feature = self.model.item_texts[inter_item]['title'] + ' ' + self.model.item_texts[inter_item]['description']
+#                 inter_id = self.tokenizer(inter_feature, return_tensors = 'pt', padding=True, truncation=True).to(self.model.device)
+#                 inter_emb = llama_model(input_ids = inter_id.input_ids, attention_mask = inter_id.attention_mask)
+#                 inter_emb = inter_emb.last_hidden_state * inter_id.attention_mask.unsqueeze(-1)
+#                 inter_emb = inter_emb.sum(dim=1) / inter_id.attention_mask.sum(dim = -1, keepdim = True)
+#                 inter_emb_list.append(inter_emb.detach())
+#             inter_embs = torch.cat(inter_emb_list, dim = 0)
+#             item_feature = self.model.item_texts[d['item']]['title'] + ' ' + self.model.item_texts[d['item']]['description']
+#             item_ids = self.tokenizer(item_feature, return_tensors = 'pt', padding=True, truncation=True).to(self.model.device)
+#             item_emb = llama_model(input_ids = item_ids.input_ids, attention_mask = item_ids.attention_mask)
+#             item_emb = item_emb.last_hidden_state * item_ids.attention_mask.unsqueeze(-1)
+#             item_emb = item_emb.sum(dim=1) / item_ids.attention_mask.sum(dim = -1, keepdim = True)
+#             item_emb = item_emb.detach()
+
+#             rqids = self.model.rqvae.get_indices(torch.cat([inter_embs, item_emb], dim = 0))
+
+#             inters_rqids = rqids.view(-1, rqids.shape[-1]).cpu().numpy().tolist()[:-1]
+#             item_rqid = rqids.view(-1, rqids.shape[-1]).cpu().numpy().tolist()[-1]
+
+#             text_rqids = {}
+#             code = ''
+#             for rqid in inters_rqids:
+#                 for k, idx in enumerate(rqid):
+#                     code = code + self.model.prefix[k].format(idx)
+#                 code = code + ', '
+#             text_rqids['inters'] = code[:-2]
+#             code = ''
+#             for k, idx in enumerate(item_rqid):
+#                 code = code + self.model.prefix[k].format(idx)
+#             text_rqids['item'] = code
+
+#             d['input_ids'] = d['input_ids'].format(inters = text_rqids['inters'])
+#             d['labels'] = d['labels'].format(inters = text_rqids['inters'], item = text_rqids['item'])
+        
+#         input_texts = [d["input_ids"] for d in data]
+#         targets = [d["labels"] for d in data]
+
+#         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": false
+    },
+    "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": false
+    },
+    "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"
+    }
+}

continue_pretrain.py

@@ -0,0 +1,125 @@
+import os
+import sys
+from typing import List
+import argparse
+
+import wandb
+import torch
+import transformers
+from transformers import LlamaForCausalLM, LlamaTokenizer, LlamaConfig
+
+from peft import (
+    TaskType,
+    LoraConfig,
+    get_peft_model,
+    get_peft_model_state_dict,
+    set_peft_model_state_dict,
+)
+
+from collator import VanillaCollator
+from rq_llama import *
+from utils import *
+
+parser = argparse.ArgumentParser(description = 'rqllama-pretrain')
+parser = parse_global_args(parser)
+parser = parse_train_args(parser)
+parser = parse_dataset_args(parser)
+parser = parse_rqvae_args(parser)
+args = parser.parse_args()
+wandb.init(config = args, reinit = True)
+
+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}
+
+train_data, valid_data = load_datasets(args)
+
+rqllama = LlamaWithRQ.from_pretrained(args.ckpt_path, torch_dtype = torch.float16, low_cpu_mem_usage = True, device_map = device_map)
+for i in range(len(args.num_emb_list)):
+    rqllama.item_rqvae.rq.vq_layers[i].initted = True
+    rqllama.user_rqvae.rq.vq_layers[i].initted = True
+
+if local_rank == 0:
+    print("token num:", len(rqllama.tokenizer))
+    print("data num:", len(train_data))
+    rqllama.tokenizer.save_pretrained(args.output_dir)
+    rqllama.config.save_pretrained(args.output_dir)
+
+if args.resume_from_checkpoint:
+    checkpoint_name = os.path.join(args.resume_from_checkpoint, "adapter_model.bin")
+    args.resume_from_checkpoint = False
+    if os.path.exists(checkpoint_name):
+        if local_rank == 0:
+            print(f"Restarting from {checkpoint_name}")
+        adapters_weights = torch.load(checkpoint_name)
+        rqllama.model = set_peft_model_state_dict(rqllama.model, adapters_weights)
+    else:
+        if local_rank == 0:
+            print(f"Checkpoint {checkpoint_name} not found")
+
+if local_rank == 0:
+    rqllama.model.print_trainable_parameters()
+
+if not ddp and torch.cuda.device_count() > 1:
+    rqllama.is_parallelizable = True
+    rqllama.model_parallel = True
+
+collator = VanillaCollator(args, rqllama.tokenizer)
+
+trainer = transformers.Trainer(
+    model = rqllama,
+    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,
+        remove_unused_columns = args.remove_unused_columns,
+    ),
+    tokenizer = rqllama.tokenizer,
+    data_collator = collator,
+)
+rqllama.config.use_cache = False
+
+if torch.__version__ >= "2" and sys.platform != "win32":
+    rqllama = torch.compile(rqllama)
+
+trainer.train(resume_from_checkpoint = args.resume_from_checkpoint)
+
+trainer.save_state()
+trainer.save_model(output_dir = args.output_dir)
+
+if local_rank == 0:
+    print('rqllama pre-train finished.')

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,1035 @@
+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.user_index_file = args.user_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]:
+                    return list(self.allowed_tokens[i])
+
+        return prefix_allowed_tokens_fn
+
+    def _process_data(self):
+        raise NotImplementedError
+
+class UserFeatDataset(BaseDataset):
+    def __init__(self, args, task = "pref2user", 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]
+
+        self._load_data()
+        self.feat_data = self._process_data()
+
+    def _load_data(self):
+        with open(os.path.join(self.data_path, self.dataset + ".user.json"), 'r') as f:
+            user_feat = json.load(f)
+            # >>> user_feat.keys()
+            # dict_keys(['user_explicit_preference', 'user_vague_intention'])
+            self.user_feat = user_feat['user_explicit_preference']
+            # >>> user_feat['0']
+            # ['The user is a passionate musician who enjoys exploring different types of musical instruments.']
+            # >>> len(user_feat)
+            # 24772
+
+    def _process_data(self):
+        feat_data = []
+        for uid in self.user_feat:
+            one_data = {}
+            one_data['user'] = uid
+
+            preference = " ".join(self.user_feat[uid])
+            preference = preference.strip().strip(".!?,;:`")
+            preference = preference.replace('{','').replace('}','')
+            one_data['preference'] = preference
+
+            feat_data.append(one_data)
+
+        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 __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]
+
+        if self.task == 'pref2user':
+            instruction = prompt['instruction'].format(preference = d['preference'])
+            input = sft_prompt.format(instruction = instruction, response = "")
+            output = sft_prompt.format(instruction = instruction, response = prompt["response"])
+            return dict(
+                input_ids = input, 
+                labels = output, 
+                inters = 'placeholder', 
+                item = 'placeholder',
+                users = 'placeholder',
+                user = d['user'], 
+                task = self.task
+            )
+        elif self.task == 'user2pref':
+            input = sft_prompt.format(instruction = prompt["instruction"], response = "")
+            response = prompt["response"].format(preference = d['preference'])
+            output = sft_prompt.format(instruction = prompt["instruction"], response = response)
+            return dict(
+                input_ids = input, 
+                labels = output, 
+                inters = 'placeholder', 
+                item = 'placeholder',
+                users = 'placeholder',
+                user = d['user'], 
+                task = self.task
+            )
+        else:
+            raise NotImplementedError
+
+class UserSearchDataset(BaseDataset):
+    def __init__(self, args, prompt_sample_num = 1, prompt_id = 0, sample_num = -1):
+        super().__init__(args)
+
+        self.prompt_sample_num = prompt_sample_num
+        self.prompt_id = prompt_id
+        self.sample_num = sample_num
+
+        self.prompts = all_prompt["usersearch"]
+
+        self._load_data()
+        self.search_data = self._process_data()
+        
+    def _load_data(self):
+        with open(os.path.join(self.data_path, self.dataset + ".inter.user.json"), 'r') as f:
+            self.user_inters = json.load(f)
+
+    def _process_data(self):
+        search_data = []
+        for iid in self.user_inters.keys():
+            users = self.user_inters[iid]
+            for i in range(1, len(users)):
+                one_data = {}
+                one_data['item'] = iid
+                one_data['user'] = str(users[i])
+                history = users[:i]
+
+                if len(history) > self.max_his_len:
+                    history = history[-self.max_his_len:]
+
+                one_data['users'] = ''
+                for user in history:
+                    one_data['users'] = one_data['users'] + str(user) + ','
+                one_data['users'] = one_data['users'][:-1]
+
+                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 __len__(self):
+        return len(self.search_data) * self.prompt_sample_num
+
+    def __getitem__(self, index):
+        idx = index // self.prompt_sample_num
+        d = self.search_data[idx]
+
+        prompt_id = random.randint(0, len(self.prompts) - 1)
+        prompt = self.prompts[prompt_id]
+
+        input = sft_prompt.format(instruction = prompt["instruction"], response = "")
+        output = sft_prompt.format(instruction = prompt["instruction"], response = prompt["response"])
+        return dict(
+            input_ids = input,
+            labels = output,
+            inters = 'placeholder',
+            item = d['item'], 
+            users = d['users'],
+            user = d['user'],
+            task = 'usersearch'
+        )
+
+# =====================================================================================================================
+# seqrec,itemsearch,inters2title,inters2description,preferenceobtain,item2index,index2item,intertitles2item,query2item
+# =====================================================================================================================
+
+# seqrec
+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"]
+
+        self._load_data()
+        
+        if self.mode == 'train':
+            self.inter_data = self._process_train_data()
+            # self.inter_data = self.inter_data[:10]
+        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.inter_data = self.inter_data[:10]
+            self._construct_valid_text()
+        elif self.mode == 'test':
+            self.inter_data = self._process_test_data()
+            # self.inter_data = self.inter_data[:10]
+        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)
+
+    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'] = str(items[i])
+                history = items[:i]
+                if self.max_his_len > 0:
+                    history = history[-self.max_his_len:]
+                one_data['inters'] = ''
+                for item in history:
+                    one_data['inters'] = one_data['inters'] + str(item) + ','
+                one_data['inters'] = one_data['inters'][:-1]
+
+                inter_data.append(one_data)
+
+        return inter_data
+    
+    def _process_valid_data(self):
+        inter_data = []
+        for uid in self.inters:
+            one_data = dict()
+            items = self.inters[uid]
+            one_data['user'] = uid
+            one_data['item'] = str(items[-2])
+            history = items[:-2]
+            if self.max_his_len > 0:
+                history = history[-self.max_his_len:]
+            one_data['inters'] = ''
+            for item in history:
+                one_data['inters'] = one_data['inters'] + str(item) + ','
+            one_data['inters'] = one_data['inters'][:-1]
+            inter_data.append(one_data)
+
+        return inter_data
+
+    def _process_test_data(self):
+        with open(self.index_file, 'r') as f:
+            self.indices = json.load(f)
+        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
+
+        with open(self.user_index_file, 'r') as f:
+            self.user_indices = json.load(f)
+        self.remapped_users = dict()
+        for uid in self.inters:
+            new_user= ''.join(self.user_indices[uid])
+            self.remapped_users[uid] = new_user
+        
+        inter_data = []
+        for uid in self.remapped_inters:
+            one_data = dict()
+            one_data['user'] = self.remapped_users[uid]
+            items = self.remapped_inters[uid]
+            one_data['item'] = items[-1]
+            history = items[:-1]
+            if self.max_his_len > 0:
+                history = history[-self.max_his_len:]
+            one_data["inters"] = self.his_sep.join(history)
+            inter_data.append(one_data)
+
+        # for uid in self.inters:
+        #     one_data = dict()
+        #     items = self.inters[uid]
+        #     one_data["item"] = str(items[-1])
+        #     history = items[:-1]
+        #     if self.max_his_len > 0:
+        #         history = history[-self.max_his_len:]
+        #     one_data['inters'] = ''
+        #     for item in history:
+        #         one_data['inters'] = one_data['inters'] + str(item) + ','
+        #     one_data['inters'] = one_data['inters'][:-1]
+        #     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 = sft_prompt.format(instruction = prompt["instruction"], response = "")
+                    output = sft_prompt.format(instruction = prompt["instruction"], response = prompt["response"])
+                    self.valid_text_data.append({
+                        "input_ids": input, 
+                        "labels": output, 
+                        "inters": d['inters'], 
+                        "item": d['item'],
+                        "users": 'placeholder',
+                        "user": d['user'],
+                        "task": 'seqrec'})
+        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 = sft_prompt.format(instruction = prompt["instruction"], response = "")
+                output = sft_prompt.format(instruction = prompt["instruction"], response = prompt["response"])
+                self.valid_text_data.append({
+                    "input_ids": input, "labels": output, 
+                    "inters": d['inters'], "item": d['item'], "users": 'placeholder', "user": d['user'], 
+                    "task": 'seqrec'})
+
+    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]
+            instruction = prompt["instruction"].format(**d)
+            response = prompt["response"].format(**d)
+            input = sft_prompt.format(instruction = instruction, response = "")
+            return dict(input_ids = input, labels = response)
+            # output = prompt["response"]
+            # return dict(input_ids = input, labels = output, inters = d['inters'], item = d['item'], task = 'seqrec')
+
+        prompt = self.prompts[prompt_id]
+
+        input = sft_prompt.format(instruction = prompt["instruction"], response = "")
+        output = sft_prompt.format(instruction = prompt["instruction"], response = prompt["response"])
+
+        return dict(input_ids = input, labels = output, inters = d['inters'], item = d['item'], user = d['user'],  task = 'seqrec', users = 'placeholder')
+
+# itemsearch & query2item
+class ItemSearchDataset(BaseDataset):
+    def __init__(self, args, mode="train", task = 'itemsearch',
+                 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.task = task.lower()
+        self.prompts = all_prompt[self.task]
+
+        self._load_data()
+        self.search_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)
+
+    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 = {}
+            one_data['user'] = uid
+            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"]
+
+            if len(inters) == 0:
+                continue
+
+            one_data["item"] = str(iid)
+
+            if self.max_his_len > 0:
+                inters = inters[-self.max_his_len:]
+            one_data["inters"] = ''            
+            for item in inters:
+                one_data["inters"] = one_data["inters"] + str(item) + ','
+            one_data["inters"] = one_data["inters"][:-1]
+
+            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"]))
+        d["explicit_preference"] = d["explicit_preference"].replace('{','').replace('}','')
+        d["user_related_intention"] = d["user_related_intention"].replace('{','').replace('}','')
+        d["item_related_intention"] = d["item_related_intention"].replace('{','').replace('}','')
+        all_querys = [d["user_related_intention"], d["item_related_intention"]]
+        d["query"] = random.choice(all_querys)
+
+        # d["query"] = d["query"].replace('{','').replace('}','')
+
+        if self.task == 'itemsearch':
+            sub_d = d.copy()
+            sub_d.pop('inters')
+            sub_d.pop('user')
+            instruction = prompt["instruction"].format(inters='{inters}', user='{user}', **sub_d)
+            input = sft_prompt.format(instruction = instruction, response = "")
+            output = sft_prompt.format(instruction = instruction, response = prompt["response"])
+            return dict(input_ids = input, labels = output, inters = d['inters'], item = d['item'], user = d['user'], task = self.task, users = 'placeholder')
+        elif self.task == 'query2item':
+            sub_d = d.copy()
+            sub_d.pop('user')
+            instruction = prompt["instruction"].format(user='{user}', **sub_d)
+            input = sft_prompt.format(instruction = instruction, response = "")
+            output = sft_prompt.format(instruction = instruction, response = prompt["response"])
+            return dict(input_ids = input, labels = output, inters = 'placeholder', item = d['item'], user = d['user'], task = self.task, users = 'placeholder')
+
+# inters2title & inters2description & intertitles2item
+class FusionSeqRecDataset(BaseDataset):
+    def __init__(self, args, mode="train", task = 'inters2title',
+                 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.task = task.lower()
+        self.prompts = all_prompt[self.task]
+
+        # load data
+        self._load_data()
+
+        # 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(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["item"] = str(items[i])
+                one_data['user'] = uid
+                one_data["title"] = self.item_feat[str(items[i])]["title"].strip().strip(".!?,;:`")
+                one_data["title"] = one_data["title"].replace('{','').replace('}','')
+                one_data["description"] = self.item_feat[str(items[i])]["description"]
+                one_data["description"] = one_data["description"].replace('{','').replace('}','')
+
+                history = items[:i]
+                if self.max_his_len > 0:
+                    history = history[-self.max_his_len:]
+                
+                one_data['inters'] = ''
+                for item in history:
+                    one_data['inters'] = one_data['inters'] + str(item) +','
+                one_data['inters'] = one_data['inters'][:-1]
+
+                inter_titles = ["\"" + self.item_feat[str(j)]["title"].strip().strip(".!?,;:`").replace('{','').replace('}','') + "\"" for j in history]
+                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"] = 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"]
+            one_data["description"] = one_data["description"].replace('{','').replace('}','')
+
+            history = items[:-2]
+            if self.max_his_len > 0:
+                history = history[-self.max_his_len:]
+            one_data['inters'] = ''
+            for item in history:
+                one_data['inters'] = one_data['inters'] + str(item) +','
+            one_data['inters'] = one_data['inters'][:-1]
+
+            inter_titles = ["\"" + self.item_feat[str(j)]["title"].strip().strip(".!?,;:`") + "\"" for j in history]
+            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"] = 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:]
+
+            one_data['inters'] = ''
+            for item in history:
+                one_data['inters'] = one_data['inters'] + str(item) +','
+            one_data['inters'] = one_data['inters'][:-1]
+
+            inter_titles = ["\"" + self.item_feat[str(j)]["title"].strip().strip(".!?,;:`") + "\"" for j in history]
+            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)
+                if self.task == 'inters2title':
+                    for prompt_id in prompt_ids:
+                        prompt = self.prompts[prompt_id]
+                        input = sft_prompt.format(instruction = prompt['instruction'], response = "")
+                        response = prompt['response'].format(title = d['title'])
+                        output = sft_prompt.format(instruction = prompt['instruction'], response = response)
+                        self.valid_text_data.append({"input_ids": input, "labels": output, 'inters': d['inters'], 'item': 'placeholder', 'task': self.task})
+                elif self.task == 'inters2description':
+                    for prompt_id in prompt_ids:
+                        prompt = self.prompts[prompt_id]
+                        input = sft_prompt.format(instruction = prompt['instruction'], response = "")
+                        response = prompt['response'].format(title = d['description'])
+                        output = sft_prompt.format(instruction = prompt['instruction'], response = response)
+                        self.valid_text_data.append({"input_ids": input, "labels": output, 'inters': d['inters'], 'item': 'placeholder', 'task': self.task})
+                elif self.task == 'intertitles2item':
+                    for prompt_id in prompt_ids:
+                        prompt = self.prompts[prompt_id]
+                        instruction = prompt['instruction'].format(inter_titles = d['inter_titles'])
+                        input = sft_prompt.format(instruction = instruction, response = "")
+                        output = sft_prompt.format(instruction = instruction, response = prompt["response"])
+                        self.valid_text_data.append({"input_ids": input, "labels": output, 'inters': 'placeholder', 'item': d['item'], 'task': self.task})
+                else:
+                    raise NotImplementedError
+        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]
+                if self.task == 'inters2title':
+                    input = sft_prompt.format(instruction = prompt['instruction'], response = "")
+                    response = prompt['response'].format(title = d['title'])
+                    output = sft_prompt.format(instruction = prompt['instruction'], response = response)
+                    self.valid_text_data.append({"input_ids": input, "labels": output, 'inters': d['inters'], 'item': 'placeholder', 'task': self.task})
+                elif self.task == 'inters2description':
+                    input = sft_prompt.format(instruction = prompt['instruction'], response = "")
+                    response = prompt['response'].format(title = d['description'])
+                    output = sft_prompt.format(instruction = prompt['instruction'], response = response)
+                    self.valid_text_data.append({"input_ids": input, "labels": output, 'inters': d['inters'], 'item': 'placeholder', 'task': self.task})
+                elif self.task == 'intertitles2item':
+                    instruction = prompt['instruction'].format(inter_titles = d['inter_titles'])
+                    input = sft_prompt.format(instruction = instruction, response = "")
+                    output = sft_prompt.format(instruction = instruction, response = prompt["response"])
+                    self.valid_text_data.append({"input_ids": input, "labels": output, 'inters': 'placeholder', 'item': d['item'], 'task': self.task})
+                else:
+                    raise NotImplementedError
+
+    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]
+
+        if self.task == 'inters2title':
+            input = sft_prompt.format(instruction = prompt['instruction'], response = "")
+            response = prompt['response'].format(title = d['title'])
+            output = sft_prompt.format(instruction = prompt['instruction'], response = response)
+            return dict(input_ids = input, labels = output, inters = d['inters'], user = d['user'], item = 'placeholder', task = self.task, users = 'placeholder')
+        elif self.task == 'inters2description':
+            input = sft_prompt.format(instruction = prompt['instruction'], response = "")
+            response = prompt['response'].format(description = d['description'])
+            output = sft_prompt.format(instruction = prompt['instruction'], response = response)
+            return dict(input_ids = input, labels = output, inters = d['inters'], user = d['user'], item = 'placeholder', task = self.task, users = 'placeholder')
+        elif self.task == 'intertitles2item':
+            instruction = prompt['instruction'].format(user = '{user}', inter_titles = d['inter_titles'])
+            input = sft_prompt.format(instruction = instruction, response = "")
+            output = sft_prompt.format(instruction = instruction, response = prompt["response"])
+            return dict(input_ids = input, labels = output, inters = 'placeholder', user = d['user'], item = d['item'], task = self.task, users = 'placeholder')
+        else:
+            raise NotImplementedError
+
+# preferenceobtain
+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.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)
+
+    def _process_data(self):
+        preference_data = []
+        user_explicit_preference = self.user_info["user_explicit_preference"]
+
+        for uid in user_explicit_preference.keys():
+            one_data = {}
+            one_data['user'] = uid
+            inters = self.inters[uid][:-3]
+            user_ep = user_explicit_preference[uid]
+
+            if self.max_his_len > 0:
+                inters = inters[-self.max_his_len:]
+            one_data['inters'] = ''
+            for item in inters:
+                one_data['inters'] = one_data['inters'] + str(item) + ','
+            one_data['inters'] = one_data['inters'][:-1]
+
+            one_data["explicit_preferences"] = user_ep
+
+            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"]))
+        d["explicit_preference"] = d["explicit_preference"].replace('{','').replace('}','')
+
+        input = sft_prompt.format(instruction = prompt["instruction"], response = "")
+        response = prompt["response"].format(**d)
+        output = sft_prompt.format(instruction = prompt["instruction"], response = response)
+        return dict(input_ids = input, labels = output, inters = d['inters'], user = d['user'], item = 'placeholder', task = 'preferenceobtain', users = 'placeholder')
+
+# item2index & index2item
+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]
+
+        self._load_data()
+        self.feat_data = self._process_data()
+
+    def _load_data(self):
+        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]
+            feat["item"] = iid
+            feat["title"] = feat["title"].strip().strip(".!?,;:`")
+            feat["title"] = feat["title"].replace('{','').replace('}','')
+            feat["description"] = feat["description"].strip().strip(".!?,;:`")
+            feat["description"] = feat["description"].replace('{','').replace('}','')
+            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]
+
+        if self.task == 'item2index':
+            instruction = prompt["instruction"].format(**d)
+            input = sft_prompt.format(instruction = instruction, response = "")
+            output = sft_prompt.format(instruction = instruction, response = prompt["response"])
+            return dict(input_ids = input, labels = output, inters = 'placeholder', user = 'placeholder', item = d['item'], task = self.task, users = 'placeholder')
+        elif self.task == 'index2item':
+            input = sft_prompt.format(instruction = prompt["instruction"], response = "")
+            response = prompt["response"].format(**d)
+            output = sft_prompt.format(instruction = prompt["instruction"], response = response)
+            return dict(input_ids = input, labels = output, inters = 'placeholder', user = 'placeholder', item = d['item'], task = self.task, users = 'placeholder')
+        else:
+            raise NotImplementedError
+
+
+
+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(os.path.join(self.data_path, self.dataset + 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_finetune.py

@@ -0,0 +1,1026 @@
+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_finetune 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.user_index_file = args.user_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 UserSearchFinetune(BaseDataset):
+    def __init__(self, args, prompt_sample_num = 1, prompt_id = 0, sample_num = -1):
+        super().__init__(args)
+
+        self.prompt_sample_num = prompt_sample_num
+        self.prompt_id = prompt_id
+        self.sample_num = sample_num
+
+        self.prompts = all_prompt["usersearch"]
+
+        self._load_data()
+        self._remap_items()
+        self.search_data = self._process_data()
+        
+    def _load_data(self):
+        with open(os.path.join(self.data_path, self.dataset + ".inter.user.json"), 'r') as f:
+            self.user_inters = json.load(f)
+        with open(self.user_index_file, 'r') as f:
+            self.user_indices = json.load(f)
+        with open(self.index_file, 'r') as f:
+            self.indices = json.load(f)
+
+    def _remap_items(self):
+        self.remapped_user_inters = dict()
+        for iid, users in self.user_inters.items():
+            new_users = ["".join(self.user_indices[str(i)]) for i in users]
+            self.remapped_user_inters[iid] = new_users
+
+    def _process_data(self):
+        search_data = []
+        for iid in self.remapped_user_inters.keys():
+            users = self.remapped_user_inters[iid]
+            for i in range(1, len(users)):
+                one_data = {}
+                one_data['item'] = self.indices[iid]
+                one_data['user'] = users[i]
+                history = users[:i]
+
+                if len(history) > self.max_his_len:
+                    history = history[-self.max_his_len:]
+
+                one_data['users'] = self.his_sep.join(history)
+
+                # one_data['users'] = ''
+                # for user in history:
+                #     one_data['users'] = one_data['users'] + str(user) + ','
+                # one_data['users'] = one_data['users'][:-1]
+
+                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 __len__(self):
+        return len(self.search_data) * self.prompt_sample_num
+
+    def __getitem__(self, index):
+        idx = index // self.prompt_sample_num
+        d = self.search_data[idx]
+
+        prompt_id = random.randint(0, len(self.prompts) - 1)
+        prompt = self.prompts[prompt_id]
+
+        instruction = prompt["instruction"].format(**d)
+        response = prompt["response"].format(**d)
+
+        input = sft_prompt.format(instruction = prompt["instruction"], response = "")
+        output = sft_prompt.format(instruction = prompt["instruction"], response = prompt["response"])
+
+        return dict(input_ids = input, labels = output)
+        # return dict(
+        #     input_ids = input,
+        #     labels = output,
+        #     inters = 'placeholder',
+        #     item = d['item'], 
+        #     users = d['users'],
+        #     user = d['user'],
+        #     task = 'usersearch'
+        # )
+
+class UserFeatFinetune(BaseDataset):
+    def __init__(self, args, task = "pref2user", 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]
+
+        self._load_data()
+        self.feat_data = self._process_data()
+
+    def _load_data(self):
+        with open(os.path.join(self.data_path, self.dataset + ".user.json"), 'r') as f:
+            user_feat = json.load(f)
+            self.user_feat = user_feat['user_explicit_preference']
+        with open(self.user_index_file, 'r') as f:
+            self.user_indices = json.load(f)
+
+    def _process_data(self):
+        feat_data = []
+        for uid in self.user_feat:
+            one_data = {}
+            one_data['user'] = self.user_indices[uid]
+
+            preference = " ".join(self.user_feat[uid])
+            preference = preference.strip().strip(".!?,;:`")
+            preference = preference.replace('{','').replace('}','')
+            one_data['preference'] = preference
+
+            feat_data.append(one_data)
+
+        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 __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]
+
+        instruction = prompt["instruction"].format(**d)
+        response = prompt["response"].format(**d)
+
+        input = sft_prompt.format(instruction = prompt["instruction"], response = "")
+        output = sft_prompt.format(instruction = prompt["instruction"], response = prompt["response"])
+
+        return dict(input_ids = input, labels = output)
+
+        # if self.task == 'pref2user':
+        #     instruction = prompt['instruction'].format(preference = d['preference'])
+        #     input = sft_prompt.format(instruction = instruction, response = "")
+        #     output = sft_prompt.format(instruction = instruction, response = prompt["response"])
+        #     return dict(
+        #         input_ids = input, 
+        #         labels = output, 
+        #         inters = 'placeholder', 
+        #         item = 'placeholder',
+        #         users = 'placeholder',
+        #         user = d['user'], 
+        #         task = self.task
+        #     )
+        # elif self.task == 'user2pref':
+        #     input = sft_prompt.format(instruction = prompt["instruction"], response = "")
+        #     response = prompt["response"].format(preference = d['preference'])
+        #     output = sft_prompt.format(instruction = prompt["instruction"], response = response)
+        #     return dict(
+        #         input_ids = input, 
+        #         labels = output, 
+        #         inters = 'placeholder', 
+        #         item = 'placeholder',
+        #         users = 'placeholder',
+        #         user = d['user'], 
+        #         task = self.task
+        #     )
+        # else:
+        #     raise NotImplementedError
+
+class SeqRecFinetune(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)
+        with open(self.user_index_file, 'r') as f:
+            self.user_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"] = self.user_indices[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["user"] = self.user_indices[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["user"] = self.user_indices[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 FusionSeqRecFinetune(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(self.user_index_file, 'r') as f:
+            self.user_indices = json.load(f)
+        # with open(os.path.join(self.data_path, self.dataset + 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["user"] = self.user_indices[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 ItemFeatFinetune(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(os.path.join(self.data_path, self.dataset + self.index_file), 'r') as f:
+        #     self.indices = 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_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 ItemSearchFinetune(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(os.path.join(self.data_path, self.dataset + self.index_file), 'r') as f:
+        #     self.indices = json.load(f)
+        with open(self.index_file, 'r') as f:
+            self.indices = json.load(f)
+        with open(self.user_index_file, 'r') as f:
+            self.user_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]
+            one_data["user"] = self.user_indices[uid]
+
+            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 PreferenceObtainFinetune(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(os.path.join(self.data_path, self.dataset + self.index_file), 'r') as f:
+        #     self.indices = json.load(f)
+        with open(self.index_file, 'r') as f:
+            self.indices = json.load(f)
+        with open(self.user_index_file, 'r') as f:
+            self.user_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 = {}
+            one_data["user"] = self.user_indices[uid]
+            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(os.path.join(self.data_path, self.dataset + 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,139 @@
+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 = os.path.join(args.root, f'{args.dataset}.item.json')
+    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', 'description'])
+    # 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):
+    print(f'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 = []
+    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()
+                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)
+
+    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=2, 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')
+    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)
+
+

data_process/amazon_user_emb.py

@@ -0,0 +1,145 @@
+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 = os.path.join(args.root, f'{args.dataset}.user.json')
+    item2feature = load_json(item2feature_path)
+
+    return item2feature
+
+def generate_text(item2feature, features):
+    item_text_list = []
+
+    for item in item2feature:
+        data = item2feature[item]
+        text = []
+
+        for i in range(len(data)):
+            meta_value = clean_text(data[i])
+            text.append(meta_value.strip())
+
+        # 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)
+    item2feature = item2feature['user_explicit_preference']
+    # load item text and clean
+    item_text_list = generate_text(item2feature)
+    # item_text_list = generate_text(item2feature, ['user_explicit_preference'])
+    # 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):
+    print(f'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 = []
+    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()
+                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)
+
+    # file = os.path.join(args.root, args.dataset + '.emb-' + args.plm_name + "-td" + ".npy")
+    np.save(args.save_path, 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=2, 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('--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)

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-finetuned.py

@@ -0,0 +1,197 @@
+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
+
+parser = argparse.ArgumentParser(description = 'rqllama-evaluate')
+parser = parse_evaluate_args(parser)
+args = parser.parse_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)
+base_model = LlamaForCausalLM.from_pretrained(args.base_model, torch_dtype=torch.float16, low_cpu_mem_usage = True, device_map = device_map)
+base_model.resize_token_embeddings(len(tokenizer))
+model = PeftModel.from_pretrained(base_model, args.ckpt_path, torch_dtype = torch.float16, device_map = device_map)
+
+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)
+if local_rank == 0:
+    print("evaluate data num:", len(test_data))
+ddp_sampler = DistributedSampler(test_data, num_replicas = world_size, rank = local_rank, drop_last = True)
+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 = 4, 
+    pin_memory = True
+)
+
+model.eval()
+
+metrics = args.metrics.split(",")
+all_prompt_results = []
+
+print('prompts:', len(prompt_ids))
+
+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_ids.shape: torch.Size([20, 101])
+            # scores.shape: torch.Size([20])
+
+            output = tokenizer.batch_decode(output_ids, skip_special_tokens = True)
+            # output.length: 20
+            '''
+            Below is an instruction that describes a task. 
+            Write a response that appropriately completes the request.\n\n
+            ### Instruction:\nThe user has interacted with items  <a-213> <b-171> <c-26> <d-74> <p-0> ,  <a-14> <b-33> <c-196> <d-121> <p-0> ,  
+            <a-213> <b-23> <c-128> <d-13> <p-8> ,  <a-1> <b-23> <c-68> <d-71> <p-1> in chronological order. 
+            Can you predict the next possible item that the user may expect?\n\n
+            ### Response: <a-9> <b-23> <c-123> <d-85> <p-2>
+            '''
+
+            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)

evaluate.py

@@ -0,0 +1,69 @@
+import math
+
+def get_topk_results(predictions, scores, targets, k, all_items=None):
+    # target: ['<a-5><b-248><c-226><d-145>']
+    results = []
+    B = len(targets)
+    predictions = [_.split("Response:")[-1] for _ in predictions]
+    predictions = [_.strip().replace(" ","") for _ in predictions]
+    # prediction: ['<a-9><b-70><c-10><d-21>', '<a-5><b-88><c-103><d-74>', '<a-29><b-70><c-36><d-113>']
+
+    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)
+
+    # result: [[0, 0, 0]]
+    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
+

fine-tune.py

@@ -0,0 +1,154 @@
+import argparse
+import os
+import sys
+from typing import List
+
+import torch
+import transformers
+from peft import PeftModel
+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
+
+import argparse
+from utils import *
+from rq_llama import *
+
+parser = argparse.ArgumentParser(description = 'rqllama-finetune')
+parser = parse_finetune_args(parser)
+args = parser.parse_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}
+
+train_data, valid_data = load_finetune_datasets(args)
+
+rqllama = LlamaWithRQ.from_pretrained(args.ckpt_path, torch_dtype = torch.float16, low_cpu_mem_usage = True, device_map = device_map)
+tokenizer = rqllama.tokenizer
+# PeftModelForCausalLM
+model = rqllama.model
+device = rqllama.device
+
+postfix = '<p-{}>'
+new_tokens = []
+new_ids = list(range(args.reindex))
+for i in new_ids:
+    new_tokens.append(postfix.format(int(i)))
+tokenizer.add_tokens(new_tokens)
+
+if local_rank == 0:
+    print("token num:", len(rqllama.tokenizer))
+    print("data num:", len(train_data))
+
+collator = Collator(args, tokenizer)
+
+# Re-index Embedding
+new_ids = torch.tensor(new_ids, dtype = torch.float16).reshape(-1,1)
+re_index_emb = torch.nn.Linear(1, model.config.hidden_size, dtype = torch.float16).to(device)
+new_embeddings = re_index_emb(new_ids.to(device))
+# PeftModelForCausalLM -> LlamaForCausalLM -> LlamaModel
+model.model.model.embed_tokens.original_module.weight.data = torch.cat([model.model.model.embed_tokens.original_module.weight.data, new_embeddings], dim = 0)
+model.model.model.embed_tokens.modules_to_save.default.weight.data = torch.cat([model.model.model.embed_tokens.modules_to_save.default.weight.data, new_embeddings], dim = 0)
+
+new_lm_head = torch.randn(args.reindex, model.config.hidden_size, requires_grad = True).to(device)
+# print('new_lm_head:',new_lm_head.requires_grad)
+# PeftModelForCausalLM -> LlamaForCausalLM
+model.model.lm_head.original_module.weight.data = torch.cat([model.model.lm_head.original_module.weight.data, new_lm_head], dim = 0)
+model.model.lm_head.modules_to_save.default.weight.data = torch.cat([model.model.lm_head.modules_to_save.default.weight.data, new_lm_head], dim = 0)
+
+model.config.vocab_size = len(tokenizer)
+
+# print(model.model.model.embed_tokens.original_module.weight.shape)
+# print(len(tokenizer))
+
+model.train()
+
+if local_rank == 0:
+    model.print_trainable_parameters()
+
+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 = 50,
+        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,
+        remove_unused_columns = args.remove_unused_columns,
+    ),
+    tokenizer = tokenizer,
+    data_collator = collator,
+)
+model.config.use_cache = False
+
+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 local_rank == 0:
+    print('rqllama fine-tune finished.')
+
+    import smtplib
+    from email.mime.text import MIMEText
+    mail_host = 'smtp.qq.com'
+    mail_code = 'ouzplpngooqndjcb'
+    sender = '1849334588@qq.com'
+    receiver = 'esperanto1949@foxmail.com'
+
+    task = '[v39: finetune twin-tower]'
+    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()

finetune.py

@@ -0,0 +1,119 @@
+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,
+        ),
+        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)

generate_embeddings.py

@@ -0,0 +1,86 @@
+import os
+import collections
+import json
+import logging
+import argparse
+import numpy as np
+import pandas as pd
+import torch
+from time import time
+from torch import optim
+from tqdm import tqdm
+from torch.utils.data import DataLoader
+from rq_llama import *
+
+def parse_args():
+    parser = argparse.ArgumentParser(description = "Index")
+    parser.add_argument("--ckpt_path", type = str, default = "", help = "")
+    parser.add_argument("--item_save_path", type = str, default = "", help = "")
+    parser.add_argument("--user_save_path", type = str, default = "", help = "")
+    parser.add_argument("--device_map", type = str, default = "1", help = "gpu or cpu")
+    return parser.parse_args()
+
+args = parse_args()
+print(args)
+device_map = {'': int(args.device_map)}
+MODEL = LlamaWithRQ.from_pretrained(args.ckpt_path, torch_dtype = torch.float16, low_cpu_mem_usage = True, device_map = device_map)
+MODEL.eval()
+device = MODEL.device
+llama = MODEL.model.get_decoder()
+tokenizer = MODEL.tokenizer
+item_texts = MODEL.item_texts
+user_texts = MODEL.user_texts
+
+all_idx = []
+all_embeddings = []
+with torch.no_grad():
+    for idx, text in tqdm(item_texts.items()):
+        item_text = text['title'] + ' ' + text['description']
+        item_ids = tokenizer(item_text, return_tensors = 'pt', padding = True, truncation = True).to(device)
+        item_emb = llama(input_ids = item_ids.input_ids, attention_mask = item_ids.attention_mask)
+        item_emb = item_emb.last_hidden_state * item_ids.attention_mask.unsqueeze(-1)
+        item_emb = item_emb.sum(dim = 1) / item_ids.attention_mask.sum(dim = -1, keepdim = True)
+
+        all_idx.append(idx)
+        all_embeddings.append(item_emb.detach().cpu().numpy().flatten().tolist())
+    
+results = {
+    'id': all_idx,
+    'emb': []
+}
+
+for emb in tqdm(all_embeddings):
+    str_emb = ''
+    for e in emb:
+        str_emb = str_emb + str(e) + ' '
+    results['emb'].append(str_emb[:-1])
+
+df = pd.DataFrame(results)
+df.to_csv(args.item_save_path, sep = '\t', header = 0, index = False)
+
+all_idx = []
+all_embeddings = []
+with torch.no_grad():
+    for idx, text in tqdm(user_texts.items()):
+        user_text = ' '.join(text)
+        user_ids = tokenizer(user_text, return_tensors = 'pt', padding = True, truncation = True).to(device)
+        user_emb = llama(input_ids = user_ids.input_ids, attention_mask = user_ids.attention_mask)
+        user_emb = user_emb.last_hidden_state * user_ids.attention_mask.unsqueeze(-1)
+        user_emb = user_emb.sum(dim = 1) / user_ids.attention_mask.sum(dim = -1, keepdim = True)
+
+        all_idx.append(idx)
+        all_embeddings.append(user_emb.detach().cpu().numpy().flatten().tolist())
+
+results = {
+    'id': all_idx,
+    'emb': []
+}
+
+for emb in tqdm(all_embeddings):
+    str_emb = ''
+    for e in emb:
+        str_emb = str_emb + str(e) + ' '
+    results['emb'].append(str_emb[:-1])
+
+df = pd.DataFrame(results)
+df.to_csv(args.user_save_path, sep = '\t', header = 0, index = False)

generate_indices.py

@@ -0,0 +1,150 @@
+import os
+import collections
+import json
+import logging
+import argparse
+import numpy as np
+import pandas as pd
+import torch
+from time import time
+from torch import optim
+from tqdm import tqdm
+from torch.utils.data import DataLoader
+
+from rq_llama import *
+from index.datasets import EmbDataset
+
+def if_collided(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("--item_data_path", type = str, default = "", help = "")
+    parser.add_argument("--user_data_path", type = str, default = "", help = "")
+    parser.add_argument("--save_path", type = str, default = "", help = "")
+    parser.add_argument("--device_map", type = str, default = "1", help = "gpu or cpu")
+    return parser.parse_args()
+
+args = parse_args()
+print(args)
+
+device_map = {'': int(args.device_map)}
+MODEL = LlamaWithRQ.from_pretrained(args.ckpt_path, torch_dtype = torch.float16, low_cpu_mem_usage = True, device_map = device_map)
+MODEL.eval()
+device = MODEL.device
+postfix = '<p-{}>'
+
+data = EmbDataset(args.item_data_path)
+data_loader = DataLoader(data, num_workers = 4, batch_size = 64, shuffle = False, pin_memory = True)
+rqvae = MODEL.item_rqvae
+prefix = MODEL.prefix
+
+index_table = {}
+all_indices = []
+all_indices_str = []
+with torch.no_grad():
+    for x in tqdm(data_loader):
+        indices = rqvae.get_indices(x.to(device), 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)))
+
+            if str(code) in index_table:
+                index_table[str(code)] += 1
+            else:
+                index_table[str(code)] = 0
+            code.append(postfix.format(index_table[str(code)]))
+
+            all_indices.append(code)
+            all_indices_str.append(str(code))
+
+all_indices = np.array(all_indices)
+all_indices_str = np.array(all_indices_str)
+
+print("All indices number: ", len(all_indices))
+print("Max number of conflicts: ", max(get_indices_count(all_indices_str).values()))
+print('Re-index number:', max(index_table.values()))
+
+all_indices_dict = {}
+for item, indices in enumerate(all_indices.tolist()):
+    all_indices_dict[item] = list(indices)
+
+reindex_dict = {'reindex': max(index_table.values())}
+
+json_path = os.path.join(args.save_path,'indices.item.json')
+with open(json_path, 'w',encoding = 'utf-8') as f:
+    json.dump(all_indices_dict, f)
+
+reindex_path = os.path.join(args.save_path,'reindex.item.json')
+with open(reindex_path, 'w',encoding = 'utf-8') as f:
+    json.dump(reindex_dict, f)
+
+data = EmbDataset(args.user_data_path)
+data_loader = DataLoader(data, num_workers = 4, batch_size = 64, shuffle = False, pin_memory = True)
+rqvae = MODEL.user_rqvae
+prefix = MODEL.user_prefix
+
+# index_table = {}
+all_indices = []
+all_indices_str = []
+with torch.no_grad():
+    for x in tqdm(data_loader):
+        indices = rqvae.get_indices(x.to(device), 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)))
+
+            # if str(code) in index_table:
+            #     index_table[str(code)] += 1
+            # else:
+            #     index_table[str(code)] = 0
+            # code.append(postfix.format(index_table[str(code)]))
+
+            all_indices.append(code)
+            all_indices_str.append(str(code))
+
+all_indices = np.array(all_indices)
+all_indices_str = np.array(all_indices_str)
+
+print("All indices number: ", len(all_indices))
+print("Max number of conflicts: ", max(get_indices_count(all_indices_str).values()))
+# print('Re-index number:', max(index_table.values()))
+
+all_indices_dict = {}
+for item, indices in enumerate(all_indices.tolist()):
+    all_indices_dict[item] = list(indices)
+
+# reindex_dict = {'reindex': max(index_table.values())}
+
+json_path = os.path.join(args.save_path,'indices.user.json')
+with open(json_path, 'w',encoding = 'utf-8') as f:
+    json.dump(all_indices_dict, f)
+
+# reindex_path = os.path.join(args.save_path,'reindex.user.json')
+# with open(reindex_path, 'w',encoding = 'utf-8') as f:
+#     json.dump(reindex_dict, f)

generate_random_indices.py

@@ -0,0 +1,198 @@
+import os
+import collections
+import json
+import logging
+import argparse
+import numpy as np
+import pandas as pd
+import torch
+from time import time
+from torch import optim
+from tqdm import tqdm
+import torch.utils.data as data
+from torch.utils.data import DataLoader
+from index.models.rqvae import RQVAE
+# from rq_llama import *
+# from index.datasets import EmbDataset
+import random
+
+class NpyDataset(data.Dataset):
+    def __init__(self, data_path):
+        self.data_path = data_path
+        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)
+
+def if_collided(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("--item_model_path", type = str, default = "", help = "")
+    parser.add_argument("--item_data_path", type = str, default = "", help = "")
+    parser.add_argument("--user_model_path", type = str, default = "", help = "")
+    parser.add_argument("--user_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()
+
+generate_args = parse_args()
+print(generate_args)
+
+device = torch.device(generate_args.device)
+
+# generate item index
+ckpt = torch.load(os.path.join(generate_args.item_model_path, 'best_collision_model.pth'), map_location = torch.device('cpu'))
+args = ckpt['args']
+state_dict = ckpt['state_dict']
+
+data = NpyDataset(generate_args.item_data_path)
+data_loader = DataLoader(data, num_workers = args.num_workers, batch_size = 64, shuffle = False, pin_memory = True)
+# 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)
+
+prefix = ["<a_{}>","<b_{}>","<c_{}>","<d_{}>","<e_{}>"]
+postfix = "<p_{}>"
+
+index_table = {}
+all_indices = []
+all_indices_str = []
+with torch.no_grad():
+    for x in tqdm(data_loader):
+        # indices = model.get_indices(x.to(device), False)
+        # indices = indices.view(-1, indices.shape[-1]).cpu().numpy()
+
+        indices = np.random.randint(0, 256, size = (64, 4), dtype = int)
+        for index in indices:
+            code = []
+            for i, ind in enumerate(index):
+                code.append(prefix[i].format(int(ind)))
+
+            if str(code) in index_table:
+                index_table[str(code)] += 1
+            else:
+                index_table[str(code)] = 0
+            code.append(postfix.format(index_table[str(code)]))
+
+            all_indices.append(code)
+            all_indices_str.append(str(code))
+
+all_indices = np.array(all_indices)
+all_indices_str = np.array(all_indices_str)
+
+print("All indices number: ", len(all_indices))
+print("Max number of conflicts: ", max(get_indices_count(all_indices_str).values()))
+print('Re-index number:', max(index_table.values()))
+
+all_indices_dict = {}
+for item, indices in enumerate(all_indices.tolist()):
+    all_indices_dict[item] = list(indices)
+reindex_dict = {'reindex': max(index_table.values())}
+
+item_index_path = os.path.join(generate_args.item_model_path, 'indices.random.item.json')
+with open(item_index_path, 'w', encoding = 'utf-8') as f:
+    json.dump(all_indices_dict, f)
+
+item_reindex_path = os.path.join(generate_args.item_model_path, 'reindex.random.item.json')
+with open(item_reindex_path, 'w', encoding = 'utf-8') as f:
+    json.dump(reindex_dict, f)
+
+# generate user index
+ckpt = torch.load(os.path.join(generate_args.user_model_path, 'best_collision_model.pth'), map_location = torch.device('cpu'))
+args = ckpt['args']
+state_dict = ckpt['state_dict']
+
+data = NpyDataset(generate_args.user_data_path)
+data_loader = DataLoader(data, num_workers = args.num_workers, batch_size = 64, shuffle = False, pin_memory = True)
+# 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)
+
+prefix = ['<z-{}>','<y-{}>','<x-{}>','<w-{}>','<v-{}>']
+
+all_indices = []
+all_indices_str = []
+with torch.no_grad():
+    for x in tqdm(data_loader):
+        # indices = rqvae.get_indices(x.to(device), False)
+        # indices = indices.view(-1, indices.shape[-1]).cpu().numpy()
+        indices = np.random.randint(0, 256, size = (64, 4), dtype = int)
+        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))
+
+all_indices = np.array(all_indices)
+all_indices_str = np.array(all_indices_str)
+
+print("All indices number: ", len(all_indices))
+print("Max number of conflicts: ", max(get_indices_count(all_indices_str).values()))
+
+all_indices_dict = {}
+for item, indices in enumerate(all_indices.tolist()):
+    all_indices_dict[item] = list(indices)
+
+json_path = os.path.join(generate_args.user_model_path, 'indices.random.user.json')
+with open(json_path, 'w', encoding = 'utf-8') as f:
+    json.dump(all_indices_dict, f)

generate_static_indices.py

@@ -0,0 +1,193 @@
+import os
+import collections
+import json
+import logging
+import argparse
+import numpy as np
+import pandas as pd
+import torch
+from time import time
+from torch import optim
+from tqdm import tqdm
+from torch.utils.data import DataLoader
+from index.models.rqvae import RQVAE
+# from rq_llama import *
+# from index.datasets import EmbDataset
+
+class NpyDataset(data.Dataset):
+    def __init__(self, data_path):
+        self.data_path = data_path
+        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)
+
+def if_collided(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("--item_model_path", type = str, default = "", help = "")
+    parser.add_argument("--item_data_path", type = str, default = "", help = "")
+    parser.add_argument("--user_model_path", type = str, default = "", help = "")
+    parser.add_argument("--user_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()
+
+generate_args = parse_args()
+print(generate_args)
+
+device = torch.device(generate_args.device)
+
+# generate item index
+ckpt = torch.load(generate_args.item_model_path, map_location = torch.device('cpu'))
+args = ckpt['args']
+state_dict = ckpt['state_dict']
+
+data = NpyDataset(generate_args.item_data_path)
+data_loader = DataLoader(data, num_workers = args.num_workers, batch_size = 64, shuffle = False, pin_memory = True)
+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)
+
+prefix = ["<a_{}>","<b_{}>","<c_{}>","<d_{}>","<e_{}>"]
+postfix = "<p_{}>"
+
+index_table = {}
+all_indices = []
+all_indices_str = []
+with torch.no_grad():
+    for x in tqdm(data_loader):
+        indices = model.get_indices(x.to(device), 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)))
+
+            if str(code) in index_table:
+                index_table[str(code)] += 1
+            else:
+                index_table[str(code)] = 0
+            code.append(postfix.format(index_table[str(code)]))
+
+            all_indices.append(code)
+            all_indices_str.append(str(code))
+
+all_indices = np.array(all_indices)
+all_indices_str = np.array(all_indices_str)
+
+print("All indices number: ", len(all_indices))
+print("Max number of conflicts: ", max(get_indices_count(all_indices_str).values()))
+print('Re-index number:', max(index_table.values()))
+
+all_indices_dict = {}
+for item, indices in enumerate(all_indices.tolist()):
+    all_indices_dict[item] = list(indices)
+reindex_dict = {'reindex': max(index_table.values())}
+
+item_index_path = os.path.join(generate_args.item_model_path, 'indices.item.json')
+with open(json_path, 'w', encoding = 'utf-8') as f:
+    json.dump(all_indices_dict, f)
+
+item_reindex_path = os.path.join(generate_args.item_model_path, 'reindex.item.json')
+with open(reindex_path, 'w', encoding = 'utf-8') as f:
+    json.dump(reindex_dict, f)
+
+# generate user index
+ckpt = torch.load(generate_args.user_model_path, map_location = torch.device('cpu'))
+args = ckpt['args']
+state_dict = ckpt['state_dict']
+
+data = NpyDataset(generate_args.user_data_path)
+data_loader = DataLoader(data, num_workers = args.num_workers, batch_size = 64, shuffle = False, pin_memory = True)
+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)
+
+prefix = ['<z-{}>','<y-{}>','<x-{}>','<w-{}>','<v-{}>']
+
+all_indices = []
+all_indices_str = []
+with torch.no_grad():
+    for x in tqdm(data_loader):
+        indices = rqvae.get_indices(x.to(device), 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))
+
+all_indices = np.array(all_indices)
+all_indices_str = np.array(all_indices_str)
+
+print("All indices number: ", len(all_indices))
+print("Max number of conflicts: ", max(get_indices_count(all_indices_str).values()))
+
+all_indices_dict = {}
+for item, indices in enumerate(all_indices.tolist()):
+    all_indices_dict[item] = list(indices)
+
+json_path = os.path.join(generate_args.user_model_path, 'indices.user.json')
+with open(json_path, 'w', encoding = 'utf-8') as f:
+    json.dump(all_indices_dict, f)

index/datasets.py

@@ -0,0 +1,27 @@
+import numpy as np
+import torch
+import torch.utils.data as data
+import pandas as pd
+from tqdm import tqdm
+
+class EmbDataset(data.Dataset):
+    def __init__(self,data_path):
+        self.data_path = data_path
+        names = ['emb']
+        usecols = [1]
+        tsv_data = pd.read_csv(data_path, sep = '\t',usecols = usecols, names = names, quotechar = None, quoting = 3)
+        features = tsv_data['emb'].values.tolist()
+        num_data = len(features)
+        for i in tqdm(range(num_data)):
+            features[i] = [float(s) for s in features[i].split(' ')]
+        self.embeddings = np.array(features, dtype = np.float16)
+        assert self.embeddings.shape[0] == num_data
+        self.dim = self.embeddings.shape[-1]
+
+    def __getitem__(self, index):
+        emb = self.embeddings[index]
+        tensor_emb = torch.tensor(emb, dtype = torch.float16)
+        return tensor_emb
+
+    def __len__(self):
+        return len(self.embeddings)

index/generate_indices.py

@@ -0,0 +1,151 @@
+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("--data_path", type = str, default = "", help = "Infer data path.")
+    parser.add_argument("--ckpt_path", type=str, default="", help="model checkpoint for infer")
+    parser.add_argument("--id_save_path", type=str, default="", help="output directory for id result")
+    parser.add_argument("--device", type=str, default="cuda:0", help="gpu or cpu")
+
+    return parser.parse_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)
+
+infer_args = parse_args()
+print('infer_args:', infer_args)
+device = torch.device(infer_args.device)
+output_file = infer_args.id_save_path
+data = EmbDataset(infer_args.data_path)
+
+ckpt = torch.load(infer_args.ckpt_path, map_location = torch.device('cpu'))
+args = ckpt["args"]
+state_dict = ckpt["state_dict"]
+
+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))
+
+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 >= 20 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(output_file, '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,84 @@
+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):
+        # print('x.shape:',x.shape)
+        x = self.encoder(x)
+        x_q, rq_loss, indices = self.rq(x,use_sk=use_sk)
+        out = self.decoder(x_q)
+        # print('out.shape:',out.shape)
+
+        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,104 @@
+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])
+            Q = torch.nan_to_num(Q, Q[torch.isfinite(Q)].min().item())
+            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,8 @@
+
+python -u main.py \
+  --num_emb_list 256 256 256 256 \
+  --sk_epsilons 0.0 0.0 0.0 0.003 \
+  --device cuda:0 \
+  --data_path /data/Games/Games.emb-llama-td.npy \
+  --batch_size 1024
+

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
+
+
+

infer.sh

@@ -0,0 +1,14 @@
+CKPT_PATH=${datain}/v-yinju/rq-llama/v6/Instruments
+
+python generate_embeddings.py \
+    --ckpt_path $CKPT_PATH \
+    --item_save_path $CKPT_PATH/embeddings.item.tsv \
+    --user_save_path $CKPT_PATH/embeddings.user.tsv \
+    --device_map 0
+
+python generate_indices.py \
+    --ckpt_path $CKPT_PATH \
+    --item_data_path $CKPT_PATH/embeddings.item.tsv \
+    --user_data_path $CKPT_PATH/embeddings.user.tsv \
+    --save_path $CKPT_PATH \
+    --device_map 0

instruments_evaluate.sh

@@ -0,0 +1,18 @@
+DATASET=Instruments
+BASE_MODEL=/datain/v-yinju/llama-7b
+DATA_PATH=/datain/v-yinju/rqvae-zzx/data
+CKPT_PATH=/datain/v-yinju/rq-llama/v11.2/Ins/finetune
+RESULTS_FILE=$CKPT_PATH/eval_result.json
+
+torchrun --nproc_per_node=8 evaluate-finetuned.py \
+    --base_model $BASE_MODEL \
+    --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 \
+    --test_task seqrec \
+    --index_file /datain/v-yinju/rq-llama/v11.2/Ins/indices.item.json \
+    --user_index_file /datain/v-yinju/rq-llama/v11.2/Ins/indices.user.json

instruments_finetune.sh

@@ -0,0 +1,33 @@
+export WANDB_MODE=disabled
+export CUDA_LAUNCH_BLOCKING=0
+
+DATASET=Instruments
+CKPT_PATH=/datain/v-yinju/rq-llama/v11/Instruments
+DATA_PATH=/datain/v-yinju/rqvae-zzx/data
+OUTPUT_DIR=$CKPT_PATH/finetune
+
+torchrun --nproc_per_node=8 fine-tune.py \
+    --ckpt_path $CKPT_PATH \
+    --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,itemsearch,preferenceobtain,item2index,index2item,fusionseqrec,usersearch,user2pref,pref2user \
+    --train_prompt_sample_num 1,1,1,1,1,1,1,1,1 \
+    --train_data_sample_num 0,0,0,0,0,0,0,0,0 \
+    --index_file $CKPT_PATH/indices.item.json \
+    --user_index_file $CKPT_PATH/indices.user.json \
+    --reindex 17 
+
+cd convert
+nohup ./convert.sh $OUTPUT_DIR >convert.log 2>&1 &
+cd ..

instruments_more_pretrain.sh

@@ -0,0 +1,32 @@
+export WANDB_MODE=disabled
+export CUDA_LAUNCH_BLOCKING=0
+
+DATASET=Instruments
+BASE_MODEL=/datain/v-yinju/llama-7b
+CKPT_PATH=/datain/v-yinju/rq-llama/v6/Instruments
+DATA_PATH=/datain/v-yinju/rqvae-zzx/data
+OUTPUT_DIR=/datain/v-yinju/rq-llama/v3-train/Instruments/more_pretrain
+
+torchrun --nproc_per_node=8 --master_port=3324 continue_pretrain.py \
+    --base_model $BASE_MODEL \
+    --ckpt_path $CKPT_PATH \
+    --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 \
+    --deepspeed ./config/ds_z2_fp16.json \
+    --dataloader_num_workers 4 \
+    --only_train_response \
+    --tasks seqrec,itemsearch,inters2title,inters2description,preferenceobtain,item2index,index2item,intertitles2item,query2item \
+    --train_prompt_sample_num 1,1,1,1,1,1,1,1,1 \
+    --train_data_sample_num 0,0,0,0,0,0,0,0,0 \
+    --fp16 &>>$OUTPUT_DIR/pretrain-log.txt
+
+cd convert
+nohup ./convert.sh $OUTPUT_DIR >convert.log 2>&1 &
+cd ..

instruments_pretrain.sh

@@ -0,0 +1,85 @@
+export WANDB_MODE=disabled
+export CUDA_LAUNCH_BLOCKING=0
+
+DATASET=Instruments
+BASE_MODEL=$datain/v-yinju/llama-7b
+ITEM_MODEL=$datain/v-yinju/rqvae-zzx/models/instruments/Apr-01-2024_01-25-11/best_collision_model.pth
+USER_MODEL=$datain/v-yinju/rqvae-zzx/models/instruments/user/Apr-23-2024_03-36-04/best_collision_model.pth
+DATA_PATH=$datain/v-yinju/rqvae-zzx/data
+OUTPUT_DIR=$datain/v-yinju/rq-llama/v11.2/Ins
+
+torchrun --nproc_per_node=8 pre-train.py \
+    --base_model $BASE_MODEL \
+    --item_model $ITEM_MODEL \
+    --user_model $USER_MODEL \
+    --output_dir $OUTPUT_DIR \
+    --dataset $DATASET \
+    --data_path $DATA_PATH \
+    --per_device_batch_size 6 \
+    --gradient_accumulation_steps 2 \
+    --learning_rate 5e-4 \
+    --epochs 4 \
+    --weight_decay 0.01 \
+    --save_and_eval_strategy epoch \
+    --deepspeed ./config/ds_z2_fp16.json \
+    --dataloader_num_workers 4 \
+    --only_train_response \
+    --tasks seqrec,itemsearch,inters2title,inters2description,preferenceobtain,item2index,index2item,intertitles2item,query2item,usersearch,user2pref,pref2user \
+    --train_prompt_sample_num 1,1,1,1,1,1,1,1,1,1,1,1 \
+    --train_data_sample_num 0,0,0,0,0,0,0,0,0,0,0,0 \
+    --index_file .index.json \
+    --user_index_file .user-index.json \
+    --fp16
+
+cd convert
+nohup ./convert.sh $OUTPUT_DIR >convert.log 2>&1 &
+cd ..
+
+CKPT_PATH=$datain/v-yinju/rq-llama/v11.2/Ins
+
+python generate_embeddings.py \
+    --ckpt_path $CKPT_PATH \
+    --item_save_path $CKPT_PATH/embeddings.item.tsv \
+    --user_save_path $CKPT_PATH/embeddings.user.tsv \
+    --device_map 0
+
+python generate_indices.py \
+    --ckpt_path $CKPT_PATH \
+    --item_data_path $CKPT_PATH/embeddings.item.tsv \
+    --user_data_path $CKPT_PATH/embeddings.user.tsv \
+    --save_path $CKPT_PATH \
+    --device_map 0
+
+# DATASET=Games
+# BASE_MODEL=/datain/v-yinju/llama-7b
+# ITEM_MODEL=/datain/v-yinju/rqvae-zzx/models/games/Apr-18-2024_01-51-46/best_collision_model.pth
+# USER_MODEL=/datain/v-yinju/rqvae-zzx/models/games/user/Jun-17-2024_18-40-36/best_collision_model.pth
+# DATA_PATH=/datain/v-yinju/rqvae-zzx/data
+# OUTPUT_DIR=/datain/v-yinju/rq-llama/v11/Games
+
+# torchrun --nproc_per_node=8 pre-train.py \
+#     --base_model $BASE_MODEL \
+#     --item_model $ITEM_MODEL \
+#     --user_model $USER_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 \
+#     --deepspeed ./config/ds_z2_fp16.json \
+#     --dataloader_num_workers 4 \
+#     --only_train_response \
+#     --tasks seqrec,itemsearch,inters2title,inters2description,preferenceobtain,item2index,index2item,intertitles2item,query2item,usersearch,user2pref,pref2user \
+#     --train_prompt_sample_num 1,1,1,1,1,1,1,1,1,1,1,1 \
+#     --train_data_sample_num 0,0,0,0,0,0,0,0,0,0,0,0 \
+#     --index_file .index.json \
+#     --user_index_file .user-index.json \
+#     --fp16
+
+# cd convert
+# nohup ./convert.sh $OUTPUT_DIR >convert.log 2>&1 &
+# cd ..

lora_finetune.py

@@ -0,0 +1,162 @@
+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,
+        ),
+        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)

pre-train.py

@@ -0,0 +1,154 @@
+import os
+import sys
+from typing import List
+import argparse
+
+import wandb
+import torch
+import transformers
+from transformers import LlamaForCausalLM, LlamaTokenizer, LlamaConfig
+
+from peft import (
+    TaskType,
+    LoraConfig,
+    get_peft_model,
+    get_peft_model_state_dict,
+    set_peft_model_state_dict,
+)
+
+from collator import VanillaCollator
+from rq_llama import *
+from utils import *
+
+parser = argparse.ArgumentParser(description = 'rqllama-pretrain')
+parser = parse_global_args(parser)
+parser = parse_train_args(parser)
+parser = parse_dataset_args(parser)
+parser = parse_rqvae_args(parser)
+args = parser.parse_args()
+wandb.init(config = args, reinit = True)
+
+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}
+
+train_data, valid_data = load_datasets(args)
+
+config = LlamaConfig.from_pretrained(args.base_model)
+config.args = vars(args)
+rqllama = LlamaWithRQ(config)
+
+ckpt = torch.load(args.item_model, map_location = torch.device('cpu'))
+state_dict = ckpt["state_dict"]
+rqllama.item_rqvae.load_state_dict(state_dict)
+for i in range(len(args.num_emb_list)):
+    rqllama.item_rqvae.rq.vq_layers[i].initted = True
+ckpt = torch.load(args.user_model, map_location = torch.device('cpu'))
+state_dict = ckpt["state_dict"]
+rqllama.user_rqvae.load_state_dict(state_dict)
+for i in range(len(args.num_emb_list)):
+    rqllama.user_rqvae.rq.vq_layers[i].initted = True
+
+if local_rank == 0:
+    print("token num:", len(rqllama.tokenizer))
+    print("data num:", len(train_data))
+    rqllama.tokenizer.save_pretrained(args.output_dir)
+    rqllama.config.save_pretrained(args.output_dir)
+
+if args.resume_from_checkpoint:
+    checkpoint_name = os.path.join(args.resume_from_checkpoint, "adapter_model.bin")
+    args.resume_from_checkpoint = False
+    if os.path.exists(checkpoint_name):
+        if local_rank == 0:
+            print(f"Restarting from {checkpoint_name}")
+        adapters_weights = torch.load(checkpoint_name)
+        rqllama.model = set_peft_model_state_dict(rqllama.model, adapters_weights)
+    else:
+        if local_rank == 0:
+            print(f"Checkpoint {checkpoint_name} not found")
+
+if local_rank == 0:
+    rqllama.model.print_trainable_parameters()
+
+if not ddp and torch.cuda.device_count() > 1:
+    rqllama.is_parallelizable = True
+    rqllama.model_parallel = True
+
+collator = VanillaCollator(args, rqllama.tokenizer)
+
+trainer = transformers.Trainer(
+    model = rqllama,
+    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,
+        remove_unused_columns = args.remove_unused_columns,
+    ),
+    tokenizer = rqllama.tokenizer,
+    data_collator = collator,
+)
+rqllama.config.use_cache = False
+
+if torch.__version__ >= "2" and sys.platform != "win32":
+    rqllama = torch.compile(rqllama)
+
+trainer.train(resume_from_checkpoint = args.resume_from_checkpoint)
+
+trainer.save_state()
+trainer.save_model(output_dir = args.output_dir)
+
+if local_rank == 0:
+    # print('rqllama pre-train finished.')
+
+    import smtplib
+    from email.mime.text import MIMEText
+    mail_host = 'smtp.qq.com'
+    mail_code = 'ouzplpngooqndjcb'
+    sender = '1849334588@qq.com'
+    receiver = 'esperanto1949@foxmail.com'
+
+    task = '[v53: pretrain tt.ins.5e-4 w/o projector]'
+    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()