12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849 |
- import sys
- import os
- import pandas as pd
- from gensim.models import KeyedVectors
- from gensim.models import Word2Vec
- from loguru import logger
- curr_path = os.path.abspath(os.path.dirname(__file__))
- project_root_path = curr_path[:curr_path.find("video_to_word") + len("video_to_word")]
- sys.path.append(project_root_path)
- from config.url_and_db import ai_word_product_engine
- logger.remove() # 删去 import logger之后自动产生的handler,不删除的话会出现重复输出的现象
- logger.add("/data/pythonProject/video_to_word/logs/train_word2vec_model.{time:YYYY-MM-DD}.log",
- rotation="00:00",
- format="{time:YYYY-MM-DD HH:mm:ss,SSS} [{process}] [{thread}] {level} {file} {line} - {message}",
- level="INFO")
- @logger.catch()
- def train_word2vec_model():
- # 1 get corpus
- sql = f"select word_split from tb_asr_result where word_split is not null"
- doc_df = pd.read_sql(sql, ai_word_product_engine)
- doc_df['word_split_lst'] = doc_df['word_split'].apply(lambda x: [word for word in x.split(' ')])
- logger.info(f"获取到 {doc_df['word_split_lst'].shape[0]} 个脚本文件.")
- # 2 train the gensim word2vec model with our own corpus
- model = Word2Vec(doc_df['word_split_lst'].values, min_count=5, vector_size=50, workers=3, window=3, sg=1)
- # 3 store just the words and their trained embeddings
- word_vectors = model.wv
- # 定时任务按天执行,直接覆盖历史模型文件
- word_vectors.save("/data/pythonProject/video_to_word/models/word2vec.wordvectors")
- logger.info(f"word2vec model 训练完成!")
- if __name__ == '__main__':
- # 0 train model
- train_word2vec_model()
- # 1 load back with memory-mapping = read-only,
- wv = KeyedVectors.load("/data/pythonProject/video_to_word/models/word2vec.wordvectors", mmap='r')
- # 2 get numpy vector of a word (for test)
- logger.info(f"the vector of 红包 is {wv['红包']}")
|