代码如下:
如果对于torchtext本身不太了解,可以参考:torchtext处理文本数据——构造dataset读取文本(学习一):torchtext处理文本数据——构造dataset读取文本(学习一)
from torchtext import data
COMMENT = data.Field()
fields = [('name', None), ('comment', COMMENT)]
train = data.TabularDataset.splits(
path='./',
train='my_data.csv',
format='csv',
skip_header=True,
fields=fields)
train = train[0]
# =========================上面都是torchtext的模板,下面开始自定义一个word2vec模型=========================
from gensim.models import Word2Vec
word2vec_model = Word2Vec([
['i', 'am', 'a', 'train', 'sentence', 'one'],
['i', 'am', 'a', 'love', 'sentence', 'two'],
['i', 'love', 'you'],
['you', 'love', 'me'],
], vector_size=256, min_count=1, window=5, sg=0) # 这里训练了一个超级差劲的模型,但是意思就是这个意思
word2vec_model.wv.save_word2vec_format('my_word2vec_format.txt') # 关键步骤,一定要写入到文件中
# =========================训练好后构建全部词表=========================
from torchtext.vocab import Vectors
vectors = Vectors(name='my_word2vec_format.txt') # 关键步骤,一定要从文件中读取到Vectors中
COMMENT.build_vocab(train, vectors=vectors)
love_index = COMMENT.vocab.stoi['love']
print(love_index) # love这个单词的下标
print("love的词向量:", COMMENT.vocab.vectors[love_index])