model.py

import torch
import torch.nn as nn
from peft import LoraConfig, TaskType, get_peft_model
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig
from transformers.modeling_outputs import SequenceClassifierOutputWithPast
from torch.cuda.amp import GradScaler, autocast
import numpy as np


class RecSys(nn.Module):
    def __init__(self, **args):
        super(RecSys, self).__init__()
        self.args = args
        self.input_dim, self.output_dim = args['input_dim'], args['output_dim']
        self.base_model = args["base_model"]
        self.device = torch.device(
            "cuda" if torch.cuda.is_available() else "cpu")

        # lora配置
        # 参照https://github.com/QwenLM/Qwen2.5/blob/main/examples/llama-factory/qwen2-7b-lora-sft.yaml
        peft_config = LoraConfig(task_type='CAUSAL_LM', target_modules=[
                                 "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"], r=16, lora_alpha=16, lora_dropout=0.05)

        config = AutoConfig.from_pretrained(
            self.args['base_model'], output_hidden_states=True)
        config.upcast_layernorm = True
        # model和tokenizer设置
        self.model = AutoModelForCausalLM.from_pretrained(self.base_model,
                                                          config=config,
                                                          # load_in_8bit=True,
                                                          # torch_dtype=torch.float32,
                                                          # local_files_only=True,
                                                          cache_dir='./root/')
        self.tokenizer = AutoTokenizer.from_pretrained(self.base_model,
                                                       # load_in_8bit=True,
                                                       # torch_dtype=torch.float32,
                                                       # local_files_only=True,
                                                       cache_dir='./root/')
        self.model.to(self.device)
        # 加载lora配置
        self.model = get_peft_model(self.model, peft_config)

        # 得到输入的id和mask
        instruct = "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\nGiven the user’s purchase history, predict next possible item to be purchased.\n\n### Input:\n"
        self.instruct_ids, self.instruct_mask = self.tokenizer(instruct,
                                                               # 获取两个值吗？
                                                               truncation=True, padding=False, return_tensors='pt', add_special_tokens=False).values()
        response = "\n### Response:\n"
        self.response_ids, self.response_mask = self.tokenizer(response,
                                                               truncation=True, padding=False, return_tensors='pt', add_special_tokens=False).values()

        self.embed_tokens = self.model.get_input_embeddings()
        # 嵌入层的设置
        self.item_embed = nn.Embedding.from_pretrained(
            self.args["item_embed"], freeze=True)  # 将SASRec的嵌入层权重加载进来，并冻结
        self.item_embed.to(self.device)
        self.embed_tokens.to(self.device)
        self.item_proj = nn.Linear(
            self.input_dim, self.model.config.hidden_size)  # 不确定是否可以得到config
        self.item_proj.to(self.device)

        self.score = nn.Linear(
            self.model.config.hidden_size, self.output_dim, bias=False)
        self.score.to(self.device)

    def predict(self, inputs, inputs_mask):
        bs = inputs.shape[0]
        # instruct_embeds = self.model.model.embed_tokens(self.instruct_ids.cuda()).expand(bs, -1, -1)
        # instruct_embeds = self.embed_tokens(self.instruct_ids.cuda()).expand(bs, -1, -1)
        instruct_embeds = self.embed_tokens(
            self.instruct_ids.to(self.device)).expand(bs, -1, -1)
        # response_embeds = self.model.model.embed_tokens(self.response_ids.cuda()).expand(bs, -1, -1)
        # response_embeds = self.embed_tokens(self.response_ids.cuda()).expand(bs, -1, -1)
        response_embeds = self.embed_tokens(
            self.response_ids.to(self.device)).expand(bs, -1, -1)
        instruct_mask = self.instruct_mask.to(self.device).expand(bs, -1)
        response_mask = self.response_mask.to(self.device).expand(bs, -1)

        inputs = self.item_proj(self.item_embed(inputs))
        # print("inputs:",inputs.shape)
        # print("instruct_embeds:",instruct_embeds)
        # print("response_embeds:",response_embeds)
        inputs = torch.cat([instruct_embeds, inputs, response_embeds], dim=1)
        attention_mask = torch.cat(
            [instruct_mask, inputs_mask, response_mask], dim=1)
        # assert attention_mask.size()[0] == inputs.size()[0] and attention_mask.size()[1] == inputs.size()[1]

        with autocast():  # 使用自动混合精度
            outputs = self.model(
                inputs_embeds=inputs, attention_mask=attention_mask, return_dict=True)

        # print("outputs:",dir(outputs))
        # hs=torch.tensor(outputs.hidden_states)
        # print("outputs.hidden_states.shape:",dir(outputs.hidden_states))
        pooled_output = outputs.hidden_states[-1]
        # print("pooled_output",pooled_output)
        # print("pooled_output.shape",pooled_output.shape)
        pooled_logits = self.score(pooled_output[:, -1])

        return outputs, pooled_logits.view(-1, self.output_dim)

    def forward(self, inputs, inputs_mask):
        outputs, pooled_logits = self.predict(inputs, inputs_mask)

        # loss = None
        # if labels is not None:
        #     loss_fct = nn.CrossEntropyLoss()
        #     loss = loss_fct(pooled_logits, labels.view(-1))

        # return SequenceClassifierOutputWithPast(
        #     loss=loss,
        #     logits=pooled_logits,
        #     past_key_values=outputs.past_key_values,
        #     hidden_states=outputs.hidden_states,
        #     attentions=outputs.attentions,
        # )
        # print("pooled_logits:",pooled_logits.shape)
        return pooled_logits
    
class SASRec(torch.nn.Module):
    def __init__(self, item_num, args):
        super(SASRec, self).__init__()

        # self.user_num = user_num
        self.item_num = item_num
        self.dev = args.device

        # TODO: loss += args.l2_emb for regularizing embedding vectors during training
        # https://stackoverflow.com/questions/42704283/adding-l1-l2-regularization-in-pytorch
        self.item_emb = torch.nn.Embedding(
            self.item_num+1, args.hidden_units, padding_idx=0)
        self.pos_emb = torch.nn.Embedding(
            args.maxlen+1, args.hidden_units, padding_idx=0)
        self.emb_dropout = torch.nn.Dropout(p=args.dropout_rate)

        self.attention_layernorms = torch.nn.ModuleList()  # to be Q for self-attention
        self.attention_layers = torch.nn.ModuleList()
        self.forward_layernorms = torch.nn.ModuleList()
        self.forward_layers = torch.nn.ModuleList()

        self.last_layernorm = torch.nn.LayerNorm(args.hidden_units, eps=1e-8)

        for _ in range(args.num_blocks):
            new_attn_layernorm = torch.nn.LayerNorm(
                args.hidden_units, eps=1e-8)
            self.attention_layernorms.append(new_attn_layernorm)

            new_attn_layer = torch.nn.MultiheadAttention(args.hidden_units,
                                                         args.num_heads,
                                                         args.dropout_rate)
            self.attention_layers.append(new_attn_layer)

            new_fwd_layernorm = torch.nn.LayerNorm(args.hidden_units, eps=1e-8)
            self.forward_layernorms.append(new_fwd_layernorm)

            new_fwd_layer = PointWiseFeedForward(
                args.hidden_units, args.dropout_rate)
            self.forward_layers.append(new_fwd_layer)

            # self.pos_sigmoid = torch.nn.Sigmoid()
            # self.neg_sigmoid = torch.nn.Sigmoid()

    def log2feats(self, log_seqs):  # TODO: fp64 and int64 as default in python, trim?
        seqs = self.item_emb(torch.LongTensor(log_seqs).to(self.dev))
        seqs *= self.item_emb.embedding_dim ** 0.5
        poss = np.tile(
            np.arange(1, log_seqs.shape[1] + 1), [log_seqs.shape[0], 1])
        # TODO: directly do tensor = torch.arange(1, xxx, device='cuda') to save extra overheads
        poss *= (log_seqs != 0)
        seqs += self.pos_emb(torch.LongTensor(poss).to(self.dev))
        seqs = self.emb_dropout(seqs)

        tl = seqs.shape[1]  # time dim len for enforce causality
        attention_mask = ~torch.tril(torch.ones(
            (tl, tl), dtype=torch.bool, device=self.dev))

        for i in range(len(self.attention_layers)):
            seqs = torch.transpose(seqs, 0, 1)
            Q = self.attention_layernorms[i](seqs)
            mha_outputs, _ = self.attention_layers[i](Q, seqs, seqs,
                                                      attn_mask=attention_mask)
            # need_weights=False) this arg do not work?
            seqs = Q + mha_outputs
            seqs = torch.transpose(seqs, 0, 1)

            seqs = self.forward_layernorms[i](seqs)
            seqs = self.forward_layers[i](seqs)

        log_feats = self.last_layernorm(seqs)  # (U, T, C) -> (U, -1, C)

        return log_feats

    def forward(self, user_ids, log_seqs, pos_seqs, neg_seqs):  # for training
        log_feats = self.log2feats(log_seqs)  # user_ids hasn't been used yet

        pos_embs = self.item_emb(torch.LongTensor(pos_seqs).to(self.dev))
        neg_embs = self.item_emb(torch.LongTensor(neg_seqs).to(self.dev))

        pos_logits = (log_feats * pos_embs).sum(dim=-1)
        neg_logits = (log_feats * neg_embs).sum(dim=-1)

        # pos_pred = self.pos_sigmoid(pos_logits)
        # neg_pred = self.neg_sigmoid(neg_logits)

        return pos_logits, neg_logits  # pos_pred, neg_pred

    def predict(self, log_seqs, item_indices):  # for inference
        log_feats = self.log2feats(log_seqs)  # user_ids hasn't been used yet

        # only use last QKV classifier, a waste
        final_feat = log_feats[:, -1, :]

        item_embs = self.item_emb(torch.LongTensor(
            item_indices).to(self.dev))  # (U, I, C)

        # print("item_embs.shape:", item_embs.shape)
        # print("item_embs:", item_embs)
        # import pickle
        # with open('item_embs.pkl', 'wb') as f:
            # pickle.dump(item_embs, f)
            # import sys
            # print("item_embs.pkl已保存至", sys.argv[0])

        logits = item_embs.matmul(final_feat.unsqueeze(-1)).squeeze(-1)

        # preds = self.pos_sigmoid(logits) # rank same item list for different users

        return logits  # preds # (U, I)


class RecClassify(nn.Module):
    def __init__(self, maxlen, dropout=0.2, class_num=2, d_model=128):
        super(RecClassify, self).__init__()
        self.maxlen = maxlen
        self.d_model = d_model
        self.embedding = nn.Embedding(60000, d_model)
        self.positional_encoding = PositionalEncoding(d_model, maxlen)
        # 定义一个 Transformer 编码器层
        self.encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=8, dim_feedforward=512)
        # 定义一个 Transformer 编码器
        self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=4)
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(d_model, class_num)

    def forward(self, inputs, masks):
        embeds = self.embedding(inputs)
        # 假设 positional_encoding 返回形状为 [seq_len, batch_size, d_model]
        embeds = self.positional_encoding(embeds)

        # 确保输入形状为 (seq_len, batch_size, d_model)
        embeds = embeds.transpose(0, 1)  # 转换为 (seq_len, batch_size, d_model)
        # print(embeds.shape)
        # print("embeds:", embeds)
        features = self.encoder(embeds, src_key_padding_mask=masks)
        # print(features.shape)
        # print("features:", features)
        features = self.dropout(features)
        pools = features.transpose(0, 1).mean(dim=1)
        # print(pools.shape)
        # print("pools:", pools)
        logits = self.fc(pools)
        return logits


class EnClassify(nn.Module):  # 带有熵、长度特征的分类器
    def __init__(self, maxlen, dropout=0.2, class_num=2, d_model=128):
        super(EnClassify, self).__init__()
        self.maxlen = maxlen
        self.d_model = d_model
        self.embedding = nn.Embedding(60000, d_model)
        self.positional_encoding = PositionalEncoding(d_model, maxlen)
        # 定义一个 Transformer 编码器层
        self.encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=8, dim_feedforward=512)
        # 定义一个 Transformer 编码器
        self.encoder = nn.TransformerEncoder(self.encoder_layer, num_layers=4)
        self.dropout = nn.Dropout(dropout)
        self.fc = nn.Linear(d_model+3, class_num)

    def forward(self, inputs, masks, eens, sens, lengths):
        embeds = self.embedding(inputs)
        # 假设 positional_encoding 返回形状为 [seq_len, batch_size, d_model]
        embeds = self.positional_encoding(embeds)

        # 确保输入形状为 (seq_len, batch_size, d_model)
        embeds = embeds.transpose(0, 1)  # 转换为 (seq_len, batch_size, d_model)
        # print(embeds.shape)
        # print("embeds:", embeds)
        features = self.encoder(embeds, src_key_padding_mask=masks)
        # print(features.shape)
        # print("features:", features)
        features = self.dropout(features)
        pools = features.transpose(0, 1).mean(dim=1)
        additional_features = torch.cat((eens, sens, lengths), dim=1)
        pools_with_additional = torch.cat((pools, additional_features), dim=1)
        # print(pools.shape)
        # print("pools:", pools)
        logits = self.fc(pools_with_additional)
        return logits


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, max_len, dropout=0.1):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(
            0, d_model, 2).float() * (-np.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        pe = pe.unsqueeze(0).transpose(0, 1)
        self.register_buffer('pe', pe)

    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)