w2v-accel 0.0.2

CUDA Word2Vec training artifact

加速Word2Vec训练的CUDA实现。

一个示例调用：

import os
from w2v_accel import train_word2vec
import time

CORPUS = "/root/word2vec_cuda/data/1bw"     # 训练语料路径
OUTPUT = os.path.abspath("vectors.bin")     # 输出模型路径

train_word2vec(
    train=CORPUS,
    output=OUTPUT,
    size=128,
    window=5,
    negative=5,
    sample=1e-3,
    min_count=5,
    iter=3,
    binary=True,
    cache_enable=True,
    debug=True,
)

print(f"Output: {OUTPUT}")
print(f"Output size: {os.path.getsize(OUTPUT) / 1024 / 1024:.2f} MB")

示例输出：

(venv) root@autodl-container-9403468337-be32b5e6:~/test# python test.py
[w2v-accel] Preparing vocab, corpus and cache ...

[w2v-accel] Training started.
[w2v-accel] [████████████████████████████████] 100.00% | alpha=0.000005 | speed=94776.2k words/s
[w2v-accel] Total training time: 24.40 seconds
[w2v-accel] Training completed. Saving models to /root/test/vectors.bin ...
[w2v-accel] Done.
Output: /root/test/vectors.bin
Output size: 275.04 MB