""" i2i - Session Word2Vec算法实现 基于用户会话序列训练Word2Vec模型，获取物品向量相似度 """ import sys import os sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))) import pandas as pd import json import argparse from datetime import datetime from collections import defaultdict from gensim.models import Word2Vec import numpy as np from db_service import create_db_connection from offline_tasks.config.offline_config import ( DB_CONFIG, OUTPUT_DIR, I2I_CONFIG, get_time_range, DEFAULT_LOOKBACK_DAYS, DEFAULT_I2I_TOP_N ) def prepare_session_data(df, session_gap_minutes=30): """ 准备会话数据 Args: df: DataFrame with columns: user_id, item_id, create_time session_gap_minutes: 会话间隔时间（分钟） Returns: List of sessions, each session is a list of item_ids """ sessions = [] # 按用户和时间排序 df = df.sort_values(['user_id', 'create_time']) # 按用户分组 for user_id, user_df in df.groupby('user_id'): user_sessions = [] current_session = [] last_time = None for _, row in user_df.iterrows(): item_id = str(row['item_id']) current_time = row['create_time'] # 判断是否需要开始新会话 if last_time is None or (current_time - last_time).total_seconds() / 60 > session_gap_minutes: if current_session: user_sessions.append(current_session) current_session = [item_id] else: current_session.append(item_id) last_time = current_time # 添加最后一个会话 if current_session: user_sessions.append(current_session) sessions.extend(user_sessions) # 过滤掉长度小于2的会话 sessions = [s for s in sessions if len(s) >= 2] return sessions def train_word2vec(sessions, config): """ 训练Word2Vec模型 Args: sessions: List of sessions config: Word2Vec配置 Returns: Word2Vec模型 """ print(f"Training Word2Vec with {len(sessions)} sessions...") model = Word2Vec( sentences=sessions, vector_size=config['vector_size'], window=config['window_size'], min_count=config['min_count'], workers=config['workers'], sg=config['sg'], epochs=config['epochs'], seed=42 ) print(f"Training completed. Vocabulary size: {len(model.wv)}") return model def generate_similarities(model, top_n=50): """ 生成物品相似度 Args: model: Word2Vec模型 top_n: Top N similar items Returns: Dict[item_id, List[Tuple(similar_item_id, score)]] """ result = {} for item_id in model.wv.index_to_key: try: similar_items = model.wv.most_similar(item_id, topn=top_n) result[item_id] = [(sim_id, float(score)) for sim_id, score in similar_items] except KeyError: continue return result def main(): parser = argparse.ArgumentParser(description='Run Session Word2Vec for i2i similarity') parser.add_argument('--window_size', type=int, default=I2I_CONFIG['session_w2v']['window_size'], help='Window size for Word2Vec') parser.add_argument('--vector_size', type=int, default=I2I_CONFIG['session_w2v']['vector_size'], help='Vector size for Word2Vec') parser.add_argument('--min_count', type=int, default=I2I_CONFIG['session_w2v']['min_count'], help='Minimum word count') parser.add_argument('--workers', type=int, default=I2I_CONFIG['session_w2v']['workers'], help='Number of workers') parser.add_argument('--epochs', type=int, default=I2I_CONFIG['session_w2v']['epochs'], help='Number of epochs') parser.add_argument('--top_n', type=int, default=DEFAULT_I2I_TOP_N, help=f'Top N similar items to output (default: {DEFAULT_I2I_TOP_N})') parser.add_argument('--lookback_days', type=int, default=DEFAULT_LOOKBACK_DAYS, help=f'Number of days to look back (default: {DEFAULT_LOOKBACK_DAYS})') parser.add_argument('--session_gap', type=int, default=30, help='Session gap in minutes') parser.add_argument('--output', type=str, default=None, help='Output file path') parser.add_argument('--save_model', action='store_true', help='Save Word2Vec model') args = parser.parse_args() # 创建数据库连接 print("Connecting to database...") engine = create_db_connection( DB_CONFIG['host'], DB_CONFIG['port'], DB_CONFIG['database'], DB_CONFIG['username'], DB_CONFIG['password'] ) # 获取时间范围 start_date, end_date = get_time_range(args.lookback_days) print(f"Fetching data from {start_date} to {end_date}...") # SQL查询 - 获取用户行为序列 sql_query = f""" SELECT se.anonymous_id AS user_id, se.item_id, se.create_time, pgs.name AS item_name FROM sensors_events se LEFT JOIN prd_goods_sku pgs ON se.item_id = pgs.id WHERE se.event IN ('click', 'contactFactory', 'addToPool', 'addToCart', 'purchase') AND se.create_time >= '{start_date}' AND se.create_time <= '{end_date}' AND se.item_id IS NOT NULL AND se.anonymous_id IS NOT NULL ORDER BY se.anonymous_id, se.create_time """ print("Executing SQL query...") df = pd.read_sql(sql_query, engine) print(f"Fetched {len(df)} records") # 转换create_time为datetime df['create_time'] = pd.to_datetime(df['create_time']) # 准备会话数据 print("Preparing session data...") sessions = prepare_session_data(df, session_gap_minutes=args.session_gap) print(f"Generated {len(sessions)} sessions") # 训练Word2Vec模型 w2v_config = { 'vector_size': args.vector_size, 'window_size': args.window_size, 'min_count': args.min_count, 'workers': args.workers, 'epochs': args.epochs, 'sg': 1 } model = train_word2vec(sessions, w2v_config) # 保存模型（可选） if args.save_model: model_path = os.path.join(OUTPUT_DIR, f'session_w2v_model_{datetime.now().strftime("%Y%m%d")}.model') model.save(model_path) print(f"Model saved to {model_path}") # 生成相似度 print("Generating similarities...") result = generate_similarities(model, top_n=args.top_n) # 创建item_id到name的映射 item_name_map = dict(zip(df['item_id'].astype(str), df.groupby('item_id')['item_name'].first())) # 输出结果 output_file = args.output or os.path.join(OUTPUT_DIR, f'i2i_session_w2v_{datetime.now().strftime("%Y%m%d")}.txt') print(f"Writing results to {output_file}...") with open(output_file, 'w', encoding='utf-8') as f: for item_id, sims in result.items(): item_name = item_name_map.get(item_id, 'Unknown') if not sims: continue # 格式：item_id \t item_name \t similar_item_id1:score1,similar_item_id2:score2,... sim_str = ','.join([f'{sim_id}:{score:.4f}' for sim_id, score in sims]) f.write(f'{item_id}\t{item_name}\t{sim_str}\n') print(f"Done! Generated i2i similarities for {len(result)} items") print(f"Output saved to: {output_file}") if __name__ == '__main__': main()