class S3RecModel(nn.Module):

def __init__(self, args):
        super(S3RecModel, self).__init__()
        # 각종 임베딩 레이어 초기화
        # 아이템 임베딩 (상품 ID -> 벡터)
        self.item_embeddings = nn.Embedding(
            args.item_size, args.hidden_size, padding_idx=0
        )
        # 속성 임베딩 (상품 속성 -> 벡터)
        self.attribute_embeddings = nn.Embedding(
            args.attribute_size, args.hidden_size, padding_idx=0
        )
        # 위치 임베딩 (시퀀스 내 위치 정보)
        self.position_embeddings = nn.Embedding(args.max_seq_length, args.hidden_size)
        
        # 인코더와 정규화 레이어들
        self.item_encoder = Encoder(args)
        self.LayerNorm = LayerNorm(args.hidden_size, eps=1e-12)
        self.dropout = nn.Dropout(args.hidden_dropout_prob)
        
        # 4가지 학습 목적을 위한 정규화 레이어들
        self.aap_norm = nn.Linear(args.hidden_size, args.hidden_size)  # Associated Attribute Prediction
        self.mip_norm = nn.Linear(args.hidden_size, args.hidden_size)  # Masked Item Prediction
        self.map_norm = nn.Linear(args.hidden_size, args.hidden_size)  # Masked Attribute Prediction
        self.sp_norm = nn.Linear(args.hidden_size, args.hidden_size)   # Segment Prediction
        
        # 이진 교차 엔트로피 손실 함수
        self.criterion = nn.BCELoss(reduction="none")
        self.apply(self.init_weights)
# 예를 들어 영화 ID가 318(쇼생크 탈출)인 경우
item_vector = model.item_embeddings(torch.tensor([318]))  # [1, hidden_size] 차원의 벡터
# 위치가 0인 경우 (시퀀스의 첫 번째 위치)
position_vector = model.position_embeddings(torch.tensor([0]))  # [1, hidden_size] 차원의 벡터

아래의 주요 학습 태스크로 나뉩니다.

AAP (Associated Attribute Prediction)

def associated_attribute_prediction(self, sequence_output, attribute_embedding):

목적: 아이템의 속성을 예측하는 태스크
예시: 영화 시퀀스가 주어졌을 때 해당 영화의 장르, 감독 등을 맞추는 것

# 예시
sequence = [아바타, 타이타닉, 터미네이터]  # 모두 제임스 캐머런 감독 영화
# -> 모델이 "이 시퀀스의 영화들은 제임스 캐머런 감독일 것이다"라고 예측하도록 학습

MIP (Masked Item Prediction)

def masked_item_prediction(self, sequence_output, target_item):

목적: 마스킹된 아이템이 무엇인지 예측
예시: BERT처럼 일부 아이템을 가리고 그게 무엇인지 맞추는 것

# 예시
원본: [아바타, 타이타닉, 터미네이터]
마스킹: [아바타, MASK, 터미네이터]
# -> 모델이 MASK 위치에 "타이타닉"이 올 것이라고 예측하도록 학습

MAP (Masked Attribute Prediction)

def masked_attribute_prediction(self, sequence_output, attribute_embedding):

목적: 마스킹된 아이템의 속성 예측
예시: 가려진 영화의 장르나 감독을 예측

시퀀스: [아바타, MASK, 터미네이터]
# -> MASK 위치의 영화는 "SF 장르일 것이다", "액션 영화일 것이다" 등을 예측

SP (Segment Prediction)

def segment_prediction(self, context, segment):

목적: 연속된 시청 세션이 서로 관련있는지 예측
예시: 사용자의 시청 기록이 하나의 일관된 세션인지 확인

세션1: [아바타, 타이타닉, 터미네이터]  # SF/액션 영화 세션
세션2: [노팅힐, 러브액츄얼리, 타이타닉]  # 로맨스 영화 세션
# -> 이 두 세션이 서로 관련있는 세션인지 아닌지 판단

사전학습 (Pretrain)

모든 태스크를 함께 사용

각 태스크의 손실값을 모두 합해서 전체 모델 학습

# trainer.py에서 받은 손실값들
aap_loss # 아이템의 속성 예측이 얼마나 정확한지
mip_loss # 마스킹된 아이템 예측이 얼마나 정확한지
map_loss # 마스킹된 아이템의 속성 예측이 얼마나 정확한지
sp_loss  # 세그먼트 예측이 얼마나 정확한지

# 이들을 합쳐서 전체 손실 계산
total_loss = aap_loss + mip_loss + map_loss + sp_loss

trainer.py의 pretrain 함수에 사용됩니다.

            aap_loss, mip_loss, map_loss, sp_loss = self.model.pretrain(
                attributes,
                masked_item_sequence,
                pos_items,
                neg_items,
                masked_segment_sequence,
                pos_segment,
                neg_segment,
            )

Finetune 메소드

# 예시 입력 데이터
batch_size = 2
seq_length = 5
input_ids = torch.tensor([
    [1, 2, 3, 4, 0],  # 첫 번째 사용자 시퀀스 (마지막은 패딩)
    [5, 6, 7, 0, 0]   # 두 번째 사용자 시퀀스 (뒤의 두 개는 패딩)
])

# 1. 어텐션 마스크 생성
attention_mask = (input_ids > 0).long()
print("Attention mask:")
print(attention_mask)
# 출력:
# tensor([[1, 1, 1, 1, 0],
#         [1, 1, 1, 0, 0]])

# 2. Subsequent 마스크 생성 (seq_length=5 예시)
subsequent_mask = torch.triu(torch.ones(1, 5, 5), diagonal=1)
subsequent_mask = (subsequent_mask == 0).long()
print("\nSubsequent mask:")
print(subsequent_mask[0])
# 출력:
# tensor([[1, 0, 0, 0, 0],
#         [1, 1, 0, 0, 0],
#         [1, 1, 1, 0, 0],
#         [1, 1, 1, 1, 0],
#         [1, 1, 1, 1, 1]])

# 3. 최종 어텐션 마스크 (두 마스크의 조합)
extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) * subsequent_mask

init_weights

# 임베딩 레이어 초기화 예시
embedding = nn.Embedding(1000, 64)  # vocab_size=1000, hidden_size=64
model.init_weights(embedding)
# -> 가중치가 평균 0, 표준편차 initializer_range인 정규분포로 초기화됨

# 선형 레이어 초기화 예시
linear = nn.Linear(64, 32)
model.init_weights(linear)
# -> 가중치는 정규분포로, bias는 0으로 초기화됨

# LayerNorm 초기화 예시
layernorm = LayerNorm(64)
model.init_weights(layernorm)
# -> weight는 1로, bias는 0으로 초기화됨
# 1. 사전학습 단계
사용자_시청_기록 = [
    "어벤져스",
    "MASK",  # 아이언맨이 가려짐
    "토르",
    "MASK",  # 캡틴아메리카가 가려짐
    "블랙팬서"
]

# 모델은 동시에:
1. 각 영화의 속성 예측 (AAP) -> "마블 영화", "액션 영화" 등
2. 가려진 영화 예측 (MIP) -> "아이언맨", "캡틴아메리카"
3. 가려진 영화의 속성 예측 (MAP) -> "슈퍼히어로물"
4. 시청 세션의 일관성 확인 (SP) -> "이것은 마블 영화 시청 세션이다"

# 2. 미세조정 단계
# 실제 추천 시스템에서:
사용자_기록 = ["어벤져스", "아이언맨", "토르"]
다음_추천 = model.predict(사용자_기록)  # "캡틴아메리카" 추천