Ver código fonte

train word2vec model

liyuyi@c-top.com.cn 3 anos atrás
pai
commit
c959a32401
3 arquivos alterados com 62 adições e 21 exclusões
  1. 6 19
      common_func.py
  2. 15 2
      config/url_and_db.py
  3. 41 0
      time_task/get_word2vec_model.py

+ 6 - 19
common_func.py

@@ -1,30 +1,17 @@
-from sqlalchemy import create_engine
-from urllib import parse
+import base64
 import json
-import numpy as np
 import logging
-from concurrent_log import ConcurrentTimedRotatingFileHandler
-from config.url_and_db import get_none_water_mark_url
-import requests
-import pandas as pd
-import base64
 
+import numpy as np
+import pandas as pd
+import requests
+from concurrent_log import ConcurrentTimedRotatingFileHandler
 
-# 创建数据库连接引擎
-def get_db_engine(db_info):
-    db_con_str = 'mysql+pymysql://%s:%s@%s:%d/%s' % \
-                 (db_info['user'],
-                  parse.quote_plus(db_info['password']),
-                  db_info['host'],
-                  db_info['port'],
-                  db_info['database'])
-    engine = create_engine(db_con_str, connect_args={'charset': 'utf8mb4'})
-    return engine
+from config.url_and_db import get_none_water_mark_url
 
 
 # insert into 替换为 replace into 存在重复唯一键时,先删除此行数据,然后插入新的数据
 def mysql_replace_into(table, conn, keys, data_iter):
-    from sqlalchemy.dialects.mysql import insert
     from sqlalchemy.ext.compiler import compiles
     from sqlalchemy.sql.expression import Insert
 

+ 15 - 2
config/url_and_db.py

@@ -1,6 +1,20 @@
 import uuid
+from urllib import parse
+
 import yaml
-from common_func import get_db_engine
+from sqlalchemy import create_engine
+
+
+# 创建数据库连接引擎
+def get_db_engine(db_info):
+    db_con_str = 'mysql+pymysql://%s:%s@%s:%d/%s' % \
+                 (db_info['user'],
+                  parse.quote_plus(db_info['password']),
+                  db_info['host'],
+                  db_info['port'],
+                  db_info['database'])
+    engine = create_engine(db_con_str, connect_args={'charset': 'utf8mb4'})
+    return engine
 
 
 with open('/data/pythonProject/video_to_word/config/config.yaml', mode='r', encoding='utf-8') as f:
@@ -18,7 +32,6 @@ else:
 jeecg_product_engine = get_db_engine(config['jeecg_boot_product_db'])
 ai_word_product_engine = get_db_engine(config['ai_word_product_db'])
 
-
 # 向腾讯云发送语音转脚本的任务请求url
 voice_to_script_task_submit_url = mac_ip_config[mac]['url'] + 'asr/task/submit'
 

+ 41 - 0
time_task/get_word2vec_model.py

@@ -0,0 +1,41 @@
+import pandas as pd
+from gensim.models import KeyedVectors
+from gensim.models import Word2Vec
+from loguru import logger
+
+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['红包']}")