Update embed_data.py

embed_data.py

@@ -2,14 +2,14 @@ from pymongo import MongoClient
 from transformers import BertTokenizer, BertModel
 import torch
 from torch.nn import Embedding
-import numpy as np
-from sklearn.metrics.pairwise import cosine_similarity
 
 # MongoDB Atlas 연결 설정
-client = MongoClient("mongodb+srv://waseoke:[email protected]/test?retryWrites=true&w=majority&tls=true&tlsAllowInvalidCertificates=true")
+client = MongoClient(
+    "mongodb+srv://waseoke:[email protected]/test?retryWrites=true&w=majority&tls=true&tlsAllowInvalidCertificates=true"
+)
 db = client["two_tower_model"]
 product_collection = db["product_tower"]
-user_collection = db['user_tower']
+user_collection = db["user_tower"]
 product_embedding_collection = db["product_embeddings"]  # 상품 임베딩을 저장할 컬렉션
 user_embedding_collection = db["user_embeddings"]  # 사용자 임베딩을 저장할 컬렉션
 
@@ -23,6 +23,7 @@ max_height = 250
 min_weight = 30
 max_weight = 200
 
+
 # 상품 타워: 데이터 임베딩
 def embed_product_data(product_data):
     # 상품명과 상세 정보 임베딩 (BERT)
@@ -45,47 +46,63 @@ def embed_product_data(product_data):
 
     # 모든 임베딩 벡터 차원 맞추기
     category_embedding = category_embedding.view(1, -1)  # 2D로 변환
-    color_embedding = color_embedding.view(1, -1)        # 2D로 변환
+    color_embedding = color_embedding.view(1, -1)  # 2D로 변환
 
     # 최종 임베딩 벡터 결합
-    combined_embedding = torch.cat((text_embedding, category_embedding, color_embedding), dim=1)
-    product_embedding = torch.nn.functional.adaptive_avg_pool1d(combined_embedding.unsqueeze(0), 512).squeeze(0)
+    combined_embedding = torch.cat(
+        (text_embedding, category_embedding, color_embedding), dim=1
+    )
+    product_embedding = torch.nn.functional.adaptive_avg_pool1d(
+        combined_embedding.unsqueeze(0), 512
+    ).squeeze(0)
 
     return product_embedding.detach().numpy()
 
+
 # 사용자 타워: 데이터 임베딩
 def embed_user_data(user_data):
     # 나이, 성별, 키, 몸무게 임베딩 (임베딩 레이어)
-    embedding_layer = Embedding(num_embeddings=100, embedding_dim=128)  # 임의로 설정된 예시 값
+    embedding_layer = Embedding(
+        num_embeddings=100, embedding_dim=128
+    )  # 임의로 설정된 예시 값
 
     # 예를 들어 성별을 'M'은 0, 'F'는 1로 인코딩
-    gender_id = 0 if user_data['gender'] == 'M' else 1
+    gender_id = 0 if user_data["gender"] == "M" else 1
 
     # 스케일링 적용
-    height = user_data['height']
-    weight = user_data['weight']
+    height = user_data["height"]
+    weight = user_data["weight"]
 
     if not (min_height <= height <= max_height):
-        raise ValueError(f"Invalid height value: {height}. Expected range: {min_height}-{max_height}")
+        raise ValueError(
+            f"Invalid height value: {height}. Expected range: {min_height}-{max_height}"
+        )
     if not (min_weight <= weight <= max_weight):
-        raise ValueError(f"Invalid weight value: {weight}. Expected range: {min_weight}-{max_weight}")
+        raise ValueError(
+            f"Invalid weight value: {weight}. Expected range: {min_weight}-{max_weight}"
+        )
 
     scaled_height = (height - min_height) * 99 // (max_height - min_height)
     scaled_weight = (weight - min_weight) * 99 // (max_weight - min_weight)
-    
-    age_embedding = embedding_layer(torch.tensor([user_data['age']])).view(1, -1)
+
+    age_embedding = embedding_layer(torch.tensor([user_data["age"]])).view(1, -1)
     gender_embedding = embedding_layer(torch.tensor([gender_id])).view(1, -1)
     height_embedding = embedding_layer(torch.tensor([scaled_height])).view(1, -1)
     weight_embedding = embedding_layer(torch.tensor([scaled_weight])).view(1, -1)
 
     # 최종 임베딩 벡터 결합
-    combined_embedding = torch.cat((age_embedding, gender_embedding, height_embedding, weight_embedding), dim=1)
-    user_embedding = torch.nn.functional.adaptive_avg_pool1d(combined_embedding.unsqueeze(0), 512).squeeze(0)
+    combined_embedding = torch.cat(
+        (age_embedding, gender_embedding, height_embedding, weight_embedding), dim=1
+    )
+    user_embedding = torch.nn.functional.adaptive_avg_pool1d(
+        combined_embedding.unsqueeze(0), 512
+    ).squeeze(0)
 
     return user_embedding.detach().numpy()
 
+
 # MongoDB Atlas에서 데이터 가져오기
-all_products = product_collection.find() # 모든 상품 데이터 가져오기
+all_products = product_collection.find()  # 모든 상품 데이터 가져오기
 all_users = user_collection.find()  # 모든 사용자 데이터 가져오기
 
 # 상품 임베딩 수행
@@ -96,10 +113,14 @@ for product_data in all_products:
     # MongoDB Atlas의 product_embeddings 컬렉션에 임베딩 저장
     product_embedding_collection.update_one(
         {"product_id": product_data["product_id"]},  # product_id 기준으로 찾기
-        {"$set": {"embedding": product_embedding.tolist()}},  # 벡터를 리스트 형태로 저장
-        upsert=True  # 기존 항목이 없으면 새로 삽입
+        {
+            "$set": {"embedding": product_embedding.tolist()}
+        },  # 벡터를 리스트 형태로 저장
+        upsert=True,  # 기존 항목이 없으면 새로 삽입
+    )
+    print(
+        f"Embedding saved to MongoDB Atlas for Product ID {product_data['product_id']}."
     )
-    print(f"Embedding saved to MongoDB Atlas for Product ID {product_data['product_id']}.")
 
 # 사용자 임베딩 수행
 for user_data in all_users:
@@ -110,9 +131,11 @@ for user_data in all_users:
         # MongoDB Atlas의 user_embeddings 컬렉션에 임베딩 저장
         user_embedding_collection.update_one(
             {"user_id": user_data["user_id"]},  # user_id 기준으로 찾기
-            {"$set": {"embedding": user_embedding.tolist()}},  # 벡터를 리스트 형태로 저장
-            upsert=True  # 기존 항목이 없으면 새로 삽입
+            {
+                "$set": {"embedding": user_embedding.tolist()}
+            },  # 벡터를 리스트 형태로 저장
+            upsert=True,  # 기존 항목이 없으면 새로 삽입
         )
         print(f"Embedding saved to MongoDB Atlas for user_id {user_data['user_id']}.")
     except ValueError as e:
-        print(f"Skipping user_id {user_data['user_id']} due to error: {e}")
+        print(f"Skipping user_id {user_data['user_id']} due to error: {e}")