从0到1实现 Word2vec 算法

欢迎关注公众号：pythonhonor，一起学习python。

前段时间，在直播行业做数据分析的朋友想应用算法提高业务指标：点击率和观看时长，问我有没有容易上手且有效的算法。这个必须有，安排。

本文做的主要工作是根据用户的历史点击行为，预测用户（user）接下来对哪些直播视频（item）感兴趣，即做到个性化推荐。本文采用 Word2vec 算法，该算法通过user行为来理解内容。它基于短时间内被浏览的 item 具有内在相似性的假设来学习 item 的embedding，可以认为与 item A 相关的其他 items 在训练数据中与 item A 是高频共现的。

1. 构建训练数据集

构建 user 数量 10000 个，item 数量 2000个。

# -*- coding: utf-8 -*-

"""
word2vec_seq.txt 格式: 
    user1    'item1,item2,item3,...,itemM'
    user2    'item1,item2,item3,...,itemN'
    user3    'item1,item2,item3,...,itemO'
    ...
    userZ    'item1,item2,item3,...,itemX'
"""

import random
from tqdm import tqdm


user_num = 10000
item_num = 2000

word2vec_seq_data = open('data/word2vec_seq.txt', 'w')

for u in tqdm(range(user_num)):
    user_id = 'user_{0}'.format(u)

    watch_item_num = random.sample(range(1, 100), 1)[0]
    item_id_list = random.sample(range(item_num), watch_item_num)
    item_id_list = ['item_{0}'.format(i) for i in item_id_list]
    item_ids = ','.join(item_id_list)

    word2vec_seq_data.write('%s\t%s\n'%(user_id, item_ids))

word2vec_seq_data.close()

2. 训练模型

# -*- coding: utf-8 -*-

import os, sys
from datetime import datetime
from gensim.models.word2vec import Word2Vec
from tqdm import tqdm


def word2vec_model(word2vec_seq_data):
    '''
    result: 
        得到 word2vec 模型
    '''
    
    # 加载 user 观看 item 数据
    user_watch_item_list = []
    word2vec_seq_data = open(word2vec_seq_data, 'r')
    for l in tqdm(word2vec_seq_data):
        user_watch_item_list.append(l.strip().split('\t')[1].split(','))
    
    word2vec_seq_data.close()
    print('Finish loading data !!!')

    
    # 训练模型
    print('Start to train word2vec model!')
    train_model_start_time = datetime.now()
    model = Word2Vec(user_watch_item_list, \
                     sg=1, \
                     size=64, \
                     window=5, \
                     min_count=3, \
                     workers=10, \
                     iter=50)
    
    train_model_end_time = datetime.now()
    print('Train model used time: {0}s.'\
          .format((train_model_end_time-train_model_start_time).seconds))

    
    # 保存模型
    model.save("data/word2vec.model")



if __name__ == '__main__':
    word2vec_seq_data = 'data/word2vec_seq.txt'
    word2vec_model(word2vec_seq_data)

此时得到一个 Word2vec 模型。

3. 获取 item 的向量数据

# -*- coding: utf-8 -*-

import os, sys
from gensim.models.word2vec import Word2Vec
from tqdm import tqdm


word2vec_model = Word2Vec.load("data/word2vec.model")

def get_item_id_vec(word2vec_seq_data):
    item_id_vec_data = open('data/item_id_vec.txt', 'w')

    item_id_dict = dict()
    with open(word2vec_seq_data) as f:
        for l in tqdm(f):
            cur_item_id_list = l.strip().split('\t')[1].split(',')
            for i in cur_item_id_list:
                
                if i in item_id_dict:
                    continue
                item_id_dict[i] = 1


    for item in tqdm(list(item_id_dict.keys())):
        try:
            item_vec = word2vec_model[item]
            item_id_vec_data.write('{0},{1}\n'.\
                                   format(item, ','.join(map(str, item_vec))))
        except:
            pass

    item_id_vec_data.close()



if __name__ == '__main__':
    word2vec_seq_data = 'data/word2vec_seq.txt'
    get_item_id_vec(word2vec_seq_data)

4. 根据user的历史行为数据，avg(item_vec_list) 得到user的向量数据

# -*- coding: utf-8 -*-

import numpy as np
from tqdm import tqdm


item_vec_dict = {}
with open('data/item_id_vec.txt') as f:
    for l in tqdm(f):
        item, item_vec = l.strip().split(',', 1)
        item_vec_dict[item] = np.array(item_vec.split(','), dtype='float')


user_vec_dict = {}
with open('data/word2vec_seq.txt') as f:
    for l in tqdm(f):
        user, items = l.strip().split('\t')
        item_list = items.split(',')
        item_vec_list = []
        for item in item_list:
            try:
                item_vec_list.append(list(item_vec_dict[item]))
            except:
                pass
            
        user_vec_dict[user] = np.mean(np.array(item_vec_list), axis=0)


        
# 保存向量数据
vec_dict_data = 'data/vec_dict.npz'
np.savez_compressed(vec_dict_data, \
                    item_vec_dict=item_vec_dict, \
                    user_vec_dict=user_vec_dict)

5. 计算相似度，即预测某 user 下一个感兴趣的 item

# -*- coding: utf-8 -*-

import numpy as np 

# load user_vec & item_vec
vec_dict = 'data/vec_dict.npz'
item_vec_dict = np.load(vec_dict, allow_pickle=True)['item_vec_dict'].item()
user_vec_dict = np.load(vec_dict, allow_pickle=True)['user_vec_dict'].item()



def cos_sim(user, item_list):
    if user=='' or item_list=='':
        return ''

    try:
        user_vec = user_vec_dict[user]
        item_vec_list = []
        for item in item_list:
            item_vec_list.append(item_vec_dict.get(item, np.zeros((1, 64))))

        cos_value = np.dot(item_vec_list, user_vec)
        cos_value_sorted = sorted(dict(zip(item_list, cos_value)).items(), \
                                  key=lambda x: x[1], \
                                  reverse=True)

        return ','.join([v[0] for v in cos_value_sorted])

    except:
        return ''


if __name__=='__main__':
    user = 'user_1'
    item_list = ['item_%s'%i for i in range(0, 10)]
    print(cos_sim(user, item_list))
    # item_8, item_1, item_6, item_4, item_0, item_3, item_5, item_9, item_7, item_2

可以看到 user_1 可能会对 item_8, item_1, item_6, item_4, item_0, item_3, item_5, item_9, item_7, item_2 比较感兴趣。