5ab1c29c
 tangwang
first commit
|
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
|
import logging
import argparse
from gensim.models import Word2Vec
import os
import multiprocessing
"""
说明
输入文件格式：
--input-file：输入的文件路径，文件内容每一行是一个空格分隔的 bid 序列（即一个随机游走结果，或者直接是一个用户行为session）。
输出文件格式：
--output-dir：输出的目录路径，保存嵌入向量和相似 bid 的文件。
生成两个文件：
bid_embeddings.txt：每个 bid 的嵌入向量，格式为 bid embedding_vector，例如：
123 0.12 0.34 0.56 ... 0.78
456 0.23 0.45 0.67 ... 0.89
bid_top_similar.txt：每个 bid 最相似的 top K 个 bid，格式为 bid similar_bid1:similarity1 similar_bid2:similarity2 ...，例如：
123 456:0.89 789:0.88 101:0.87 ...
456 123:0.89 678:0.85 234:0.84 ...
命令行参数：
--embedding-size：指定 Word2Vec 嵌入向量的维度，默认为 128。
--window：Word2Vec 模型的窗口大小，控制词的上下文范围，默认为 5。
--min-count：忽略词频低于该值的 bid，默认是 1，即不忽略任何 bid。
--workers：并行计算的线程数量，默认为 4。
--top-k：每个 bid 输出的最相似的 top K bid，默认是 200。
执行示例
假设：
输入文件路径是 input_sentences.txt
输出目录是 output/
那么可以使用以下命令：
bash
python word2vec_bid_similarity.py --input-file input_sentences.txt --output-dir output/ --embedding-size 128 --top-k 200
依赖项
请确保安装了以下依赖项：
bash
pip install gensim
"""
def train_word2vec(input_file, output_dir, embedding_size=128, window=5, min_count=1, workers=None, top_k=200, epochs=5):
"""
训练Word2Vec模型，并保存每个bid的embedding及top K相似的bid。
:param input_file: 句子文件路径
:param output_dir: 输出文件的目录路径
:param embedding_size: 嵌入维度大小
:param window: Word2Vec中的窗口大小
:param min_count: Word2Vec中忽略频次低于min_count的词
:param workers: 使用的线程数，如果为None，则设置为cpu_count-2
:param top_k: 每个bid的最相似bid的数量
:param epochs: 训练的epoch数量
"""
# 如果未设置workers，默认使用CPU核心数-2
if workers is None:
workers = max(1, multiprocessing.cpu_count() - 2)
# 检查输出目录是否存在，不存在则创建
if not os.path.exists(output_dir):
os.makedirs(output_dir)
logging.info(f"Loading sentences from {input_file}")
# 读取输入文件，格式是每行一个bid序列
sentences = []
with open(input_file, 'r') as f:
for line in f:
sentences.append(line.strip().split())
# 训练Word2Vec模型
logging.info(f"Training Word2Vec model with embedding size {embedding_size}, window {window}, epochs {epochs}, workers {workers}")
model = Word2Vec(sentences, vector_size=embedding_size, window=window, min_count=min_count, workers=workers, epochs=epochs)
# 保存每个bid的embedding
embedding_file = os.path.join(output_dir, "bid_embeddings.txt")
logging.info(f"Saving embeddings to {embedding_file}")
with open(embedding_file, 'w') as f_out:
for bid in model.wv.index_to_key:
vector = model.wv[bid]
f_out.write(f"{bid}\t{','.join(map(str, vector))}\n")
# 保存每个bid的top K相似bid
similar_file = os.path.join(output_dir, "bid_top_similar.txt")
logging.info(f"Saving top {top_k} similar bids for each bid to {similar_file}")
with open(similar_file, 'w') as f_out:
for bid in model.wv.index_to_key:
similar_bids = model.wv.most_similar(bid, topn=top_k)
similar_bids_str = ','.join([f"{similar_bid[0]}:{round(similar_bid[1], 4)}" for similar_bid in similar_bids])
f_out.write(f"{bid}\t{similar_bids_str}\n")
logging.info("Process completed successfully.")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Train Word2Vec model and calculate bid similarity.")
parser.add_argument('--input-file', type=str, required=True, help="Path to the input sentence file")
parser.add_argument('--output-dir', type=str, required=True, help="Directory to save output embeddings and similarity results")
parser.add_argument('--embedding-size', type=int, default=128, help="Size of the bid embedding vectors (default: 128)")
parser.add_argument('--window', type=int, default=5, help="Window size for Word2Vec (default: 5)")
parser.add_argument('--min-count', type=int, default=1, help="Minimum frequency of bids to be considered (default: 1)")
parser.add_argument('--workers', type=int, default=None, help="Number of workers (default: cpu_count-2)")
parser.add_argument('--top-k', type=int, default=200, help="Number of top similar bids to output (default: 200)")
parser.add_argument('--epochs', type=int, default=5, help="Number of epochs to train the model (default: 5)")
args = parser.parse_args()
# 初始化日志
logging.basicConfig(format='%(asctime)s - %(levelname)s - %(message)s', level=logging.INFO)
# 执行训练和保存结果
train_word2vec(input_file=args.input_file, output_dir=args.output_dir,
embedding_size=args.embedding_size, window=args.window,
min_count=args.min_count, workers=args.workers, top_k=args.top_k, epochs=args.epochs)
# python w2v.py --input-file input_sentences.txt --output-dir output/ --embedding-size 128 --top-k 200 --epochs 10
|