[NLP] Word2Vec- Skipgram 코드

예린·2025년 12월 17일

NLP

목록 보기
4/4

학습 데이터 구축을 위해 필요한 것들

주어지는 것

A collection of documents (str objects)

해야할 것

  1. 문서를 단어 단위로 쪼개기 → tokens_list

    각 문서 문자열을 단어 단위로 나눈다.

  2. Generate a list of tokens for each document- 토큰 목록 생성 → vocab

    전체 문서에서 중복 없는 단어 집합을 만든다. (어휘집)

  3. Create mappings between tokens and their indices → word2idx, idx2word

    각 단어에 고유한 정수 인덱스를 부여한다.

  4. Generate a list of token indices for each document → indexed_docs

    문서의 각 단어를 숫자로 변환

  5. (center word, context word) pair 생성

    윈도우 크기를 기준으로 중심 단어 주변의 단어 쌍 생성

  6. Unigram Distribution 이용하여 negative words 포함한 최종 데이터 셋 준비

모델을 학습 시킬 데이터의 one sample은 one center word index, one context word index, negative words indices (k개)로 구성된다.

1. Library Import

import random
from collections import Counter
from typing import List, Tuple
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn as nn
from datasets import load_dataset
from sklearn.decomposition import PCA
from torch.utils.data import DataLoader, Dataset

2. Hyperparameters

# hyper parameters

SEED = 42 # 난수 고정을 위한 시드
EMBED_DIM = 100 # 임베딩 벡터의 차원 수
WINDOW_SIZE = 2 # 윈도우 크기 m
NEGATIVE_SAMPLES = 5 # negative samping할 때 몇개 뽑을 건지 k
BATCH_SIZE = 512 # 한 번의 학습에서 몇 개의 샘플을 동시에 처리할 것인지 !
EPOCHS = 5 # 전체 corpus를 몇번 학습할건지
LR = 0.002 # learning rate
NUM_DOCUMENTS = 100 # 읽을 문서 수
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") # gpu있으면 gpu 아니면 cpu

SEED : 난수 고정할 때 사용한다. 같은 시드를 사용하면 매번 같은 결과 재현 가능하다.

EMBED_DIM : 각 단어를 표현할 벡터의 차원 수이다. 너무 작으면 정보 손실, 너무 크면 과적합 위험이 있다.

BATCH_SIZE : 한 번의 학습(= gradient update)에서 몇 개의 샘플을 동시에 처리할지 결정한다.

EPOCHS : 전체 데이터셋을 몇 번 반복해서 학습할지를 결정한다. 여기선 모든 문서를 5번 반복한다.

NUM_DOCUMENTS : 전체 코퍼스 중 상위 100개 문서만 불러와 학습하도록 제한한다.

3. Setting Random Seeds Across Libraries

def set_seed(seed: int):
	random.seed(seed)
	np.random.seed(seed)
	torch.manual_seed(seed)
	torch.cuda.manual_seed(seed)
	torch.cuda.manual_seed_all(seed)
	torch.use_deterministic_algorithms(True)
	torch.backends.cudnn.deterministic = True
	torch.backends.cudnn.benchmark = False
	
set_seed(SEED)

4. 문서를 단어 단위로 쪼개기

IMDB 영화 리뷰 데이터셋을 불러와서 텍스트를 tokenization하고 그 결과를 Word2Vec 모델이 학습할 수 있도록 단일 토큰 시퀀스로 flatten한다.

def read_corpus() -> List[List[str]]:
	
	imdb_dataset = load_dataset("stanfordnlp/imdb")
	
	files: List[str] = imdb_dataset["train"]["text"][:NUM_DOCUMENTS]
	print(f"files[0][:200] = {files[0][:200]}")
	
	return [[w.lower() for w in f.split()] for f in files]
	
tokens_list = read_corpus()
print(f"#(documents): {len(tokens_list)}")
print("The number of tokens in the 1st document:", len(tokens_list[0]))
print(f"imdb_corpus[0][:20] = {tokens_list[0][:20]}")

# 단일 토큰 시퀀스로 flatten
all_tokens: List[str] = [tok for doc in tokens_list for tok in doc]

[[w.lower() for w in f.split()] for f in files]

각 문서(f)를 공백 기준으로 분리하고 각 단어 w를 모두 소문자로 변환한 뒤 리스트로 만든다. 결과적으로 문서 단위의 토큰 리스트를 반환한다.

함수 read_corpus()의 결과 예시: List[List[str]]

[
  ["this", "movie", "was", "amazing", "and", "funny", ...],
  ["the", "plot", "was", "boring", "and", "predictable", ...],
  ...
]

2차원 리스트를 1차원 리스트로 변환하여 모든 문서의 단어들을 하나의 긴 시퀀스로 이어 붙인다.

예시:

all_tokens = ["this", "movie", "was", "amazing", "and", "funny",, "the", "plot", "was", "boring", "and", "predictable",]

5. Processing Token Indices

# 문서별 토큰 리스트 ex: [["i", "love", "nlp"], ["nlp", "is", "fun"]]
# 모든 문서를 하나로 flatten한 전체 단어 리스트 예: ["i", "love", "nlp", ... ]

def build_vocab(tokens_list: List[List[str]], tokens: List[str]) -> Tuple[dict, dict, List[List[int]]]:
	word2idx = {} # 단어 → 인덱스 매핑
	idx2word = {} # 인덱스 → 단어 매핑
	i = 0
	
	# 아직 등록되지 않은 단어를 사전에 추가하고 각 단어에 고유한 인덱스를 부여
	for token in tokens:
		if token not in word2idx:
			word2idx[token] = i
			idx2word[i] = token
			i += 1
		
	# 각 문서 안의 단어(token)를 해당 단어의 인덱스(word2idx[token])로 치환
	indexed_docs: List[List[int]] = [[word2idx[token] for token in tokens] for tokens in tokens_list]
	
	return word2idx, idx2word, indexed_docs

word2idx, idx2word, indexed_docs = build_vocab(tokens_list, all_tokens)
vocab_size = len(word2idx)
print(f"어휘 수: {vocab_size}")

반환값 예시:

word2idx = {'i': 0, 'love': 1, 'nlp': 2, 'is': 3, 'fun': 4}
idx2word = {0: 'i', 1: 'love', 2: 'nlp', 3: 'is', 4: 'fun'}

indexed_docs = [
  [0, 1, 2],
  [2, 3, 4]
]

6. (center word, context word) pair 생성

문서 내 단어들의 순서를 보고 중심 단어와 주변 단어의 관계를 학습 데이터로 변환하는 과정이다.

def build_skipgram_pairs_docs(indexed_docs: List[List[int]], window_size: int) -> List[Tuple[int, int]]:
	
	pairs: List[Tuple[int, int]] = [] # (center, context) 단어 인덱스 쌍들의 리스트
	
	# 문서 단위 반복
	for doc in indexed_docs:
		n = len(doc)
		
		if n <= 1:
			continue
		
		# center word 선택
		for i, center in enumerate(doc):
		
			# 윈도우 범위 계산
			left = max(0, i - window_size)
			right = min(n - 1, i + window_size)
			
			# context word 쌍 만들기
			for j in range(left, right + 1):
				
				if j == i:
					continue
					
				context = doc[j]
				pairs.append((center, context))
				
	return pairs

출력값 예시:

[(center=1, context=0), (center=1, context=2), (center=2, context=1), (center=2, context=3)]

7. Unigram Distribution

def make_unigram_probs(indexed_docs: List[List[int]], vocab_size: int, power: float = 0.75) -> np.ndarray:

	freqs = np.zeros(vocab_size, dtype=np.float64) # 각 단어의 등장 횟수를 저장할 배열
	
	for doc in indexed_docs:
		for i in doc:
			# 각 단어 등장 횟수 세기
			freqs[i] += 1
			
	# 확률 분포 계산
	probs = freqs ** power
	probs /= probs.sum()
	
	return probs
	
probs = make_unigram_probs(indexed_docs, vocab_size)

8. Dataset 생성

이 클래스는 PyTorch의 Dataset을 상속받아서 학습 시 매번 getitem으로 (center, positive, negatives) 샘플을 반환한다.

class SkipGramDataset(Dataset):

	def __init__(self, pairs: List[Tuple[int, int]], probs: np.ndarray, vocab_size: int, k: int):
		self.pairs = pairs    # (center, context) 쌍 목록
		self.probs = probs    # negative sampling 확률분포 P(w)
		self.vocab_size = vocab_size # 전체 단어 수
		self.k = k    # negative sample 개수
	
	def __len__(self):
		return len(self.pairs) # 데이터셋 길이
	
	def __getitem__(self, idx): # 샘플 하나 반환
		center, pos = self.pairs[idx] # 학습 데이터 선택
		negs = np.random.choice(self.vocab_size, size=self.k, replace=True, p=self.probs) # Negative 샘플링
		
		# Tensor 변환
		return (
		torch.tensor(center),
		torch.tensor(pos),
		torch.tensor(negs),
		)
	
dataset = SkipGramDataset(pairs, probs, vocab_size, NEGATIVE_SAMPLES)
print(f"The first example: {dataset[0]}")

출력 예시:

(tensor(42), tensor(105), tensor([ 8, 217,  3, 94, 501]))

9. SkipGram with Negative Sampling 모델

데이터를 입력받아 Loss function을 계산해주는 실제 신경망 모듈

class SkipGramNS(nn.Module):

	def __init__(self, vocab_size: int, embed_dim: int):
		super().__init__()
		self.in_embed = nn.Embedding(vocab_size, embed_dim) # center word의 임베딩 행렬
		self.out_embed = nn.Embedding(vocab_size, embed_dim) # context word의 임베딩 행렬
		nn.init.uniform_(self.in_embed.weight, a=-0.5/embed_dim, b=0.5/embed_dim) # 가중치 초기화
		nn.init.uniform_(self.out_embed.weight, a=-0.5/embed_dim, b=0.5/embed_dim)
		
    # 실제 연산 과정
    # input 전부 index 형태로
	def forward(self, center, pos, neg):
	    # (B), (D), (K) 표기는 텐서의 차원을 간단히 표현한 것
		# B: mini-batch size, K: negative samples, D: self.embed_dim
		# center: (B), pos: (B), neg: (B, K)
		
		cent = self.in_embed(center) # (B, D)
		pos = self.out_embed(pos) # (B, D)
		neg = self.out_embed(neg) # (B, K, D)
		
		# 중심-주변 단어 내적
		pos_mul = (cent * pos) # (B, D)
		pos_inner = pos_mul.sum(dim=-1) # (B)
		
		# center word랑 negative words랑 내적
        # Broadcasting 위해 !!
        center_2= center.reshape(-1, 1 ,self.embed_dim) # (B,1,D)
        
		neg_mul = cent_2 * neg # (B, K, D)
		neg_inner = neg_mul.sum(dim=-1) # (B, K)
		
		# 손실함수 계산
		pos_loss = - nn.functional.logsigmoid(pos_inner) # (B)
		neg_loss = - nn.functional.logsigmoid(-neg_inner).sum(dim=-1) # (B)
		
		# 각 샘플별 손실을 더한 후 평균
        # Batch 차원을 하나의 scalar 값으로 !!
		return (pos_loss + neg_loss).mean() 

10. train

model = SkipGramNS(vocab_size, EMBED_DIM).to(DEVICE) # SkipGarm 모델
opt = torch.optim.Adam(model.parameters(), lr=LR) # 옵티마이저

def train():

    # Batch 데이터 로드
    loader = DataLoader(
        dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        drop_last=True
    )
     
		
    for epoch in range(1, EPOCHS + 1):
        model.train()
        total_loss = 0.0

        for center, pos, neg in loader:
            center = center.to(DEVICE)
            pos = pos.to(DEVICE)
            neg = neg.to(DEVICE)

            opt.zero_grad()
            loss = model(center, pos, neg)
            loss.backward()
            opt.step()

            total_loss += loss.item()

        avg_loss = total_loss / len(loader)
        print(f"[Epoch {epoch}] mean loss: {avg_loss:.4f}")

빠른 학습을 위해 epoch를 5로 제한했지만 추가적인 하이퍼파라미터 튜닝을 통해 성능 개선이 가능할 것! :)

0개의 댓글