BrandCultivation/models/item2vec/item2vec.py

import joblib
from models.item2vec import Item2VecDataProcess
from gensim.models import Word2Vec
class Item2Vec:
    def __init__(self, city_uuid):
        self._load_data(city_uuid)
    
    def _load_data(self, city_uuid):
        """加载特征sentence"""
        data_processor = Item2VecDataProcess(city_uuid)
        self._sentences = data_processor.generate_sentence()
        
    def train(self):
        self._model = Word2Vec(
            self._sentences,
            vector_size=64,
            window=4,
            min_count=1,
            sg=1, # skip-gram
            workers=4,
            epochs=20
        )
        
    def save_model(self, model_path):
        joblib.dump(self._model, model_path)
        
        
if __name__ == "__main__":
    city_uuid = "00000000000000000000000011445301"
    model = Item2Vec(city_uuid)
    print("开始训练Item2Vec...")
    model.train()