utils3.py

import os
import re
import math
import torch
import random
import pickle
import json
import datetime
from rouge import rouge
from bleu import compute_bleu


def rouge_score(references, generated):
    """both are a list of strings"""
    score = rouge(generated, references)
    rouge_s = {k: (v * 100) for (k, v) in score.items()}
    '''
    "rouge_1/f_score": rouge_1_f,
    "rouge_1/r_score": rouge_1_r,
    "rouge_1/p_score": rouge_1_p,
    "rouge_2/f_score": rouge_2_f,
    "rouge_2/r_score": rouge_2_r,
    "rouge_2/p_score": rouge_2_p,
    "rouge_l/f_score": rouge_l_f,
    "rouge_l/r_score": rouge_l_r,
    "rouge_l/p_score": rouge_l_p,
    '''
    return rouge_s


def bleu_score(references, generated, n_gram=4, smooth=False):
    """a list of lists of tokens"""
    formatted_ref = [[ref] for ref in references]
    bleu_s, _, _, _, _, _ = compute_bleu(formatted_ref, generated, n_gram, smooth)
    return bleu_s * 100


def two_seq_same(sa, sb):
    if len(sa) != len(sb):
        return False
    for (wa, wb) in zip(sa, sb):
        if wa != wb:
            return False
    return True


def unique_sentence_percent(sequence_batch):
    unique_seq = []
    for seq in sequence_batch:
        count = 0
        for uni_seq in unique_seq:
            if two_seq_same(seq, uni_seq):
                count += 1
                break
        if count == 0:
            unique_seq.append(seq)

    return len(unique_seq) / len(sequence_batch), len(unique_seq)


def feature_detect(seq_batch, feature_set):
    feature_batch = []
    for ids in seq_batch:
        feature_list = []
        for i in ids:
            if i in feature_set:
                feature_list.append(i)
        feature_batch.append(set(feature_list))

    return feature_batch


def feature_matching_ratio(feature_batch, test_feature):
    count = 0
    for (fea_set, fea) in zip(feature_batch, test_feature):
        if fea in fea_set:
            count += 1

    return count / len(feature_batch)


def feature_coverage_ratio(feature_batch, feature_set):
    features = set()
    for fb in feature_batch:
        features = features | fb

    return len(features) / len(feature_set)


def feature_diversity(feature_batch):
    list_len = len(feature_batch)

    total_count = 0
    for i, x in enumerate(feature_batch):
        for j in range(i + 1, list_len):
            y = feature_batch[j]
            total_count += len(x & y)

    denominator = list_len * (list_len - 1) / 2
    return total_count / denominator


def mean_absolute_error(predicted, max_r, min_r, mae=True):
    total = 0
    for (r, p) in predicted:
        if p > max_r:
            p = max_r
        if p < min_r:
            p = min_r

        sub = p - r
        if mae:
            total += abs(sub)
        else:
            total += sub ** 2

    return total / len(predicted)


def root_mean_square_error(predicted, max_r, min_r):
    mse = mean_absolute_error(predicted, max_r, min_r, False)
    return math.sqrt(mse)


class EntityDictionary:
    def __init__(self):
        self.idx2entity = []
        self.entity2idx = {}

    def add_entity(self, e):
        if e not in self.entity2idx:
            self.entity2idx[e] = len(self.idx2entity)
            self.idx2entity.append(e)

    def __len__(self):
        return len(self.idx2entity)

class DataLoader0:
    def __init__(self, data_path, data_path_price_brand, index_dir, tokenizer, seq_len):
        self.user_dict = EntityDictionary() # 上面定义的class EntityDictionary类，
        self.item_dict = EntityDictionary()
        self.brand_dict = {}
        self.price_dict = {}
        self.price_brand(data_path_price_brand) # 调用下面的price_brand生成上面字典
        self.max_rating = float('-inf')
        self.min_rating = float('inf')
        self.initialize(data_path) # 调用下面initialize函数
        self.feature_set = set()
        self.tokenizer = tokenizer
        self.seq_len = seq_len
        self.train, self.valid, self.test, self.user2feature, self.item2feature = self.load_data(data_path, index_dir) # 调用下面的load_data函数

    def price_brand(self,data_path_price_brand):
        assert os.path.exists(data_path_price_brand)
        with open(data_path_price_brand, 'r', encoding='utf-8') as fr:
            item_data = json.load(fr)
  
        for item in item_data:
            item_id = item['item']
            #print(item)
            if 'brand' in item:
                item_brand = item['brand'] # 注意item.json中有缺失brand的情况,有的item没有brand信息
                self.brand_dict[item_id] = item_brand 
            else:
                self.brand_dict[item_id] = item['categories'][0][-1]

            if 'price' in item:
                item_price = item['price']
                self.price_dict[item_id] = str(round(item_price)) # 这里要注意是float浮点数！！，不是字符串, tokenizer只接受字符串和整数integer
            else:
                self.price_dict[item_id] = '0'
    
    def initialize(self, data_path):
        assert os.path.exists(data_path)
        reviews = pickle.load(open(data_path, 'rb'))
        for review in reviews:
            self.user_dict.add_entity(review['user'])
            self.item_dict.add_entity(review['item'])
            rating = review['rating'] # 提取评分，算最大和最小
            if self.max_rating < rating:
                self.max_rating = rating
            if self.min_rating > rating:
                self.min_rating = rating

    def load_data(self, data_path, index_dir):
        data = []
        reviews = pickle.load(open(data_path, 'rb'))
        for review in reviews: # reviews为列表，每条review为一个用户评论，
            (fea, adj, tem, sco) = review['template']  # (feature, adjective, sentence, sentiment)
            tokens = self.tokenizer(tem)['input_ids'] # ！！把一句话转换为每个单词的id编号，word和id一对一的关系
            # 如tokenizer('or high quality skirt')['input_ids'] 变为 [273, 1029, 3081, 23967]
            text = self.tokenizer.decode(tokens[:self.seq_len]) #keep seq_len tokens at most ,裁剪最大长度的句子，decode解码为文本句子 不是id
            data.append({'user': self.user_dict.entity2idx[review['user']],
                         'item': self.item_dict.entity2idx[review['item']],
                         'rating': review['rating'],
                         'price': 'The consumption level of this user is {}'.format(self.price_dict[review['item']]), # 用户背景写成一段话，用.format()
                         'brand': 'The user likes {}'.format(self.brand_dict[review['item']]), # 用户爱好写成一段话
                         'text': text, # 裁剪后的最大长度的句子
                         'feature': fea}) # 当前i产品特征
            self.feature_set.add(fea) # 添加到feature_set特征集合里面

        train_index, valid_index, test_index = self.load_index(index_dir) # 下面定义的
        train, valid, test = [], [], []
        user2feature, item2feature = {}, {} # 记录每个用户曾购买过的产品特征， 记录每个产品特征
        for idx in train_index:
            review = data[idx] # 从整个数据集中按训练集的index，提取数据
            train.append(review)
            u = review['user']
            i = review['item']
            f = review['feature']
            if u in user2feature:
                user2feature[u].append(f)  # 记录每个用户曾购买过的产品特征
            else:
                user2feature[u] = [f]
            if i in item2feature:
                item2feature[i].append(f)  # 记录每个产品特征
            else:
                item2feature[i] = [f]
        for idx in valid_index:
            valid.append(data[idx])
        for idx in test_index:
            test.append(data[idx])
        return train, valid, test, user2feature, item2feature

    def load_index(self, index_dir):
        assert os.path.exists(index_dir)
        with open(os.path.join(index_dir, 'train.index'), 'r') as f:
            train_index = [int(x) for x in f.readline().split(' ')]
        with open(os.path.join(index_dir, 'validation.index'), 'r') as f:
            valid_index = [int(x) for x in f.readline().split(' ')]
        with open(os.path.join(index_dir, 'test.index'), 'r') as f:
            test_index = [int(x) for x in f.readline().split(' ')]
        return train_index, valid_index, test_index
    

class DataLoader:
    def __init__(self, data_path, data_path_price_brand, data_path_user, index_dir, tokenizer, seq_len):
        self.user_dict = EntityDictionary() # 上面定义的class EntityDictionary类，
        self.item_dict = EntityDictionary()
        self.brand_dict = {} # 爱好
        self.user_review_cnt_dict = {} # 背景
        self.useful_dict = {} # 背景
        self.funny_dict = {} # 背景
        self.cool_dict = {} # 背景
        self.likes(data_path_price_brand) # 调用下面的price_brand生成上面字典 
        self.backgrounds(data_path_user) 
        self.max_rating = float('-inf')
        self.min_rating = float('inf')
        self.initialize(data_path) # 调用下面initialize函数
        self.feature_set = set()
        self.tokenizer = tokenizer
        self.seq_len = seq_len
        self.train, self.valid, self.test, self.user2feature, self.item2feature = self.load_data(data_path, index_dir) # 调用下面的load_data函数

    def likes(self,data_path_price_brand):
        assert os.path.exists(data_path_price_brand)
        with open(data_path_price_brand, 'r', encoding='utf-8') as fr:
            item_data = json.load(fr)
  
        for item in item_data:
            item_id = item['item']
            #print(item)
            if ',' not in item['categories']: # item只有一个categories，没有逗号分隔
                item_brand = item['categories']
                self.brand_dict[item_id] = item_brand 
            else:
                item_brand = item['categories'].split(', ')[-1] # 取最后一个categories，注意item.json中有缺失brand的情况,有的item没有brand信息
                self.brand_dict[item_id] = item_brand 
            #if 'categories' in item: # 爱好, 店铺的类型
                #item_brand = item['categories'] # 注意item.json中有缺失brand的情况,有的item没有brand信息
                #self.brand_dict[item_id] = item_brand 
            #else:
                #self.brand_dict[item_id] = item['categories'][0][-1]

            #if 'price' in item:
                #item_price = item['price']
                #self.price_dict[item_id] = str(round(item_price)) # 这里要注意是float浮点数！！，不是字符串, tokenizer只接受字符串和整数integer
            #else:
                #self.price_dict[item_id] = '0'
                
    def backgrounds(self,data_path_user): # 
        assert os.path.exists(data_path_user)
        with open(data_path_user, 'r', encoding='utf-8') as fr:
            user_data = json.load(fr)
  
        for user in user_data:
            user_id = user['user']
            #print(item)
            #if 'review_count' in user: # 用户的背景信息用，评论数量，粉丝数量，评论有用 有趣 酷数量'useful''funny' 'cool'
            user_review_cnt = user['review_count'] # 注意user.json中没有缺失review_count的情况,
            self.user_review_cnt_dict[user_id] = user_review_cnt

            useful = user['useful'] # 注意user.json中没有缺失review_count的情况,
            self.useful_dict[user_id] = useful
            
            funny = user['funny'] # 注意user.json中没有缺失review_count的情况,
            self.funny_dict[user_id] = funny
            
            cool = user['cool'] # 注意user.json中没有缺失review_count的情况,
            self.cool_dict[user_id] = cool
            #if 'price' in item:
                #item_price = item['price']
                #self.price_dict[item_id] = str(round(item_price)) # 这里要注意是float浮点数！！，不是字符串, tokenizer只接受字符串和整数integer
            
    
    def initialize(self, data_path):
        assert os.path.exists(data_path)
        reviews = pickle.load(open(data_path, 'rb'))
        for review in reviews:
            self.user_dict.add_entity(review['user'])
            self.item_dict.add_entity(review['item'])
            rating = review['rating'] # 提取评分，算最大和最小
            if self.max_rating < rating:
                self.max_rating = rating
            if self.min_rating > rating:
                self.min_rating = rating

    def load_data(self, data_path, index_dir):
        data = []
        reviews = pickle.load(open(data_path, 'rb'))
        for review in reviews: # reviews为列表，每条review为一个用户评论，
            (fea, adj, tem, sco) = review['template']  # (feature, adjective, sentence, sentiment)
            tokens = self.tokenizer(tem)['input_ids'] # ！！把一句话转换为每个单词的id编号，word和id一对一的关系
            # 如tokenizer('or high quality skirt')['input_ids'] 变为 [273, 1029, 3081, 23967]
            text = self.tokenizer.decode(tokens[:self.seq_len]) #keep seq_len tokens at most ,裁剪最大长度的句子，decode解码为文本句子 不是id
            data.append({'user': self.user_dict.entity2idx[review['user']],
                         'item': self.item_dict.entity2idx[review['item']],
                         'rating': review['rating'],
                         'price': 'This user has released {} reviews and the useful socre is {} and the funny score is {} and the cool score is {}'.format(self.user_review_cnt_dict[review['user']], self.useful_dict[review['user']], self.funny_dict[review['user']], self.cool_dict[review['user']]), # 用户背景写成一段话，用.format()
                         'brand': 'This user likes {}'.format(self.brand_dict[review['item']]), #用户爱好写成一段话
                         'text': text, # 裁剪后的最大长度的句子
                         'feature': fea}) # 当前i产品特征
            self.feature_set.add(fea) # 添加到feature_set特征集合里面

        train_index, valid_index, test_index = self.load_index(index_dir)
        train, valid, test = [], [], []
        user2feature, item2feature = {}, {} # 记录每个用户曾购买过的产品特征， 记录每个产品特征
        for idx in train_index:
            review = data[idx] # 从整个数据集中按训练集的index，提取数据
            train.append(review)
            u = review['user']
            i = review['item']
            f = review['feature']
            if u in user2feature:
                user2feature[u].append(f)  # 记录每个用户曾购买过的产品特征
            else:
                user2feature[u] = [f]
            if i in item2feature:
                item2feature[i].append(f)  # 记录每个产品特征
            else:
                item2feature[i] = [f]
        for idx in valid_index:
            valid.append(data[idx])
        for idx in test_index:
            test.append(data[idx])
        return train, valid, test, user2feature, item2feature

    def load_index(self, index_dir):
        assert os.path.exists(index_dir)
        with open(os.path.join(index_dir, 'train.index'), 'r') as f:
            train_index = [int(x) for x in f.readline().split(' ')]
        with open(os.path.join(index_dir, 'validation.index'), 'r') as f:
            valid_index = [int(x) for x in f.readline().split(' ')]
        with open(os.path.join(index_dir, 'test.index'), 'r') as f:
            test_index = [int(x) for x in f.readline().split(' ')]
        return train_index, valid_index, test_index


class Batchify:
    def __init__(self, data, tokenizer, bos, eos, batch_size=128, shuffle=False):
        u, i, r, t, self.feature, price, brand = [], [], [], [], [], [], [] # r为u对i的评分；t为加了开头bos结尾eos的句子
        for x in data: # 一般data代的就是args.train训练集数据。构造prompt
            u.append(x['user']) 
            i.append(x['item'])
            r.append(x['rating'])
            t.append('{} {} {}'.format(bos, x['text'], eos)) # bos = '<bos>' ； eos = '<eos>' ；x['text']为裁剪后最大长度的句子,
            # t句子变为 '<bos> or high quality skirt <eos>'
            self.feature.append(x['feature']) # 当前i产品特征
            price.append('{} {} {}'.format(bos, x['price'], eos)) # 用户背景，文本数据
            brand.append('{} {} {}'.format(bos, x['brand'], eos)) # brand是文本数据 用户爱好

        # 现在 u, i, r, t, self.feature, price, brand 已经经过一次for循环，添加了数据
        encoded_inputs = tokenizer(t, padding=True, return_tensors='pt') # 这一步是转为id了，
        # t = '<bos> or high quality skirt <eos>' 
        # 变为{'input_ids': tensor([[50257,   393,  1029,  3081, 23967,   220, 50258]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1]])}
        self.seq = encoded_inputs['input_ids'].contiguous()
        self.mask = encoded_inputs['attention_mask'].contiguous()
        
        encoded_inputs_brand = tokenizer(brand, padding=True, return_tensors='pt') # 这一步是转为id了，用户爱好
        self.seq_brand = encoded_inputs_brand['input_ids'].contiguous() # 注意是seq_brand！
        self.mask_brand = encoded_inputs_brand['attention_mask'].contiguous() # 注意是mask_brand！
        
        encoded_inputs_price = tokenizer(price, padding=True, return_tensors='pt') # 这一步是转为id了，用户背景
        self.seq_price = encoded_inputs_price['input_ids'].contiguous() # 注意是seq_brand！
        self.mask_price = encoded_inputs_price['attention_mask'].contiguous() # 注意是mask_brand！
        
        self.user = torch.tensor(u, dtype=torch.int64).contiguous()
        self.item = torch.tensor(i, dtype=torch.int64).contiguous()
        self.rating = torch.tensor(r, dtype=torch.float).contiguous() # 这里评分就是浮点数
        
        self.shuffle = shuffle
        self.batch_size = batch_size
        self.sample_num = len(data) # 训练集data的总样本数，
        self.index_list = list(range(self.sample_num))
        self.total_step = int(math.ceil(self.sample_num / self.batch_size)) # math.ceil（x）返回大于或等于 x 的的最小整数
        self.step = 0

    def next_batch(self):
        if self.step == self.total_step:
            self.step = 0
            if self.shuffle:
                random.shuffle(self.index_list)

        start = self.step * self.batch_size
        offset = min(start + self.batch_size, self.sample_num)
        self.step += 1
        index = self.index_list[start:offset]
        user = self.user[index]  # (batch_size,)
        item = self.item[index]
        rating = self.rating[index]
        seq = self.seq[index]  # (batch_size, seq_len)
        mask = self.mask[index]
        
        seq_price = self.seq_price[index]
        mask_price = self.mask_price[index]
        seq_brand = self.seq_brand[index]
        mask_brand = self.mask_brand[index]
        
        return user, item, rating, seq, mask, seq_price, mask_price, seq_brand, mask_brand


class Batchify2: # 这个是用在discrete prompt learning里面
    def __init__(self, data, user2feature, item2feature, tokenizer, bos, eos, seq_len, batch_size=128, shuffle=False):
        t, self.feature, features = [], [], []
        for x in data:
            ufea = set(user2feature[x['user']])
            ifea = set(item2feature[x['item']])
            intersection = ufea & ifea
            difference = ufea | ifea - intersection
            features.append(' '.join(list(intersection) + list(difference)))
            t.append('{} {} {}'.format(bos, x['text'], eos))
            self.feature.append(x['feature'])

        encoded_inputs = tokenizer(t, padding=True, return_tensors='pt')
        self.seq = encoded_inputs['input_ids'].contiguous()
        self.mask = encoded_inputs['attention_mask'].contiguous()
        encoded_features = tokenizer(features, padding=True, return_tensors='pt')
        self.prompt = encoded_features['input_ids'][:, :seq_len].contiguous()
        self.shuffle = shuffle
        self.batch_size = batch_size
        self.sample_num = len(data)
        self.index_list = list(range(self.sample_num))
        self.total_step = int(math.ceil(self.sample_num / self.batch_size))
        self.step = 0

    def next_batch(self):
        if self.step == self.total_step:
            self.step = 0
            if self.shuffle:
                random.shuffle(self.index_list)

        start = self.step * self.batch_size
        offset = min(start + self.batch_size, self.sample_num)
        self.step += 1
        index = self.index_list[start:offset]
        seq = self.seq[index]  # (batch_size, seq_len)
        mask = self.mask[index]
        prompt = self.prompt[index]
        return seq, mask, prompt


def now_time():
    return '[' + datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f') + ']: '


def postprocessing(string):
    '''
    adopted from https://github.com/yoonkim/CNN_sentence/blob/master/process_data.py
    '''
    string = re.sub('\'s', ' \'s', string)
    string = re.sub('\'m', ' \'m', string)
    string = re.sub('\'ve', ' \'ve', string)
    string = re.sub('n\'t', ' n\'t', string)
    string = re.sub('\'re', ' \'re', string)
    string = re.sub('\'d', ' \'d', string)
    string = re.sub('\'ll', ' \'ll', string)
    string = re.sub('\(', ' ( ', string)
    string = re.sub('\)', ' ) ', string)
    string = re.sub(',+', ' , ', string)
    string = re.sub(':+', ' , ', string)
    string = re.sub(';+', ' . ', string)
    string = re.sub('\.+', ' . ', string)
    string = re.sub('!+', ' ! ', string)
    string = re.sub('\?+', ' ? ', string)
    string = re.sub(' +', ' ', string).strip()
    return string


def ids2tokens(ids, tokenizer, eos):
    text = tokenizer.decode(ids)
    text = postprocessing(text)  # process punctuations: "good!" -> "good !"
    tokens = []
    for token in text.split():
        if token == eos:
            break
        tokens.append(token)
    return tokens