Spaces:
Sleeping
Sleeping
File size: 5,773 Bytes
eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e eb7ba37 e6dc15e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 |
from pymongo import MongoClient
from transformers import BertTokenizer, BertModel
import torch
from torch.nn import Embedding
# MongoDB Atlas 연결 설정
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"]
product_embedding_collection = db["product_embeddings"] # 상품 임베딩을 저장할 컬렉션
user_embedding_collection = db["user_embeddings"] # 사용자 임베딩을 저장할 컬렉션
# Hugging Face의 한국어 BERT 모델 및 토크나이저 로드 (예: klue/bert-base)
tokenizer = BertTokenizer.from_pretrained("klue/bert-base")
model = BertModel.from_pretrained("klue/bert-base")
# Height와 Weight 스케일링에 필요한 값 설정
min_height = 50
max_height = 250
min_weight = 30
max_weight = 200
# 상품 타워: 데이터 임베딩
def embed_product_data(product_data):
# 상품명과 상세 정보 임베딩 (BERT)
text = product_data.get("title", "") + " " + product_data.get("description", "")
inputs = tokenizer(
text, return_tensors="pt", truncation=True, padding=True, max_length=128
)
outputs = model(**inputs)
text_embedding = outputs.last_hidden_state.mean(dim=1) # 평균 풀링으로 벡터화
# 카테고리 및 색상 정보 임베딩 (임베딩 레이어)
category_embedding_layer = Embedding(num_embeddings=50, embedding_dim=16)
color_embedding_layer = Embedding(num_embeddings=20, embedding_dim=8)
category_id = product_data.get("category_id", 0) # 카테고리 ID, 기본값 0
color_id = product_data.get("color_id", 0) # 색상 ID, 기본값 0
category_embedding = category_embedding_layer(torch.tensor([category_id]))
color_embedding = color_embedding_layer(torch.tensor([color_id]))
# 모든 임베딩 벡터 차원 맞추기
category_embedding = category_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)
return product_embedding.detach().numpy()
# 사용자 타워: 데이터 임베딩
def embed_user_data(user_data):
# 나이, 성별, 키, 몸무게 임베딩 (임베딩 레이어)
embedding_layer = Embedding(
num_embeddings=100, embedding_dim=128
) # 임의로 설정된 예시 값
# 예를 들어 성별을 'M'은 0, 'F'는 1로 인코딩
gender_id = 0 if user_data["gender"] == "M" else 1
# 스케일링 적용
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}"
)
if not (min_weight <= 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)
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)
return user_embedding.detach().numpy()
# MongoDB Atlas에서 데이터 가져오기
all_products = product_collection.find() # 모든 상품 데이터 가져오기
all_users = user_collection.find() # 모든 사용자 데이터 가져오기
# 상품 임베딩 수행
for product_data in all_products:
product_embedding = embed_product_data(product_data)
print(f"Product ID {product_data['product_id']} Embedding: {product_embedding}")
# MongoDB Atlas의 product_embeddings 컬렉션에 임베딩 저장
product_embedding_collection.update_one(
{"product_id": product_data["product_id"]}, # product_id 기준으로 찾기
{
"$set": {"embedding": product_embedding.tolist()}
}, # 벡터를 리스트 형태로 저장
upsert=True, # 기존 항목이 없으면 새로 삽입
)
print(
f"Embedding saved to MongoDB Atlas for Product ID {product_data['product_id']}."
)
# 사용자 임베딩 수행
for user_data in all_users:
try:
user_embedding = embed_user_data(user_data)
print(f"User ID {user_data['user_id']} Embedding:", user_embedding)
# MongoDB Atlas의 user_embeddings 컬렉션에 임베딩 저장
user_embedding_collection.update_one(
{"user_id": user_data["user_id"]}, # user_id 기준으로 찾기
{
"$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}")
|