임베딩 모델에 새로운 단어 가르치기

from sentence_transformers import SentenceTransformer
import torch

# 베이스 모델 로드
model_name = 'sentence-transformers/all-MiniLM-L6-v2'
model = SentenceTransformer(model_name)

# 1. 내부 tokenizer와 모델 얻기
# 수정하려면 하위의 'transformers' 모델에 접근해야 한다
tokenizer = model.tokenizer
transformer_model = model.auto_model

print(f"Original vocabulary size: {len(tokenizer)}")

# 2. 새 단어 정의
# 중요: 모델이 'cased'인지 'uncased'인지 먼저 확인할 것!
# 이 모델('all-MiniLM-L6-v2')은 uncased라서 전부 소문자로 바꾼다.
# 소문자 버전만 추가하면 된다.
# 'cased' 모델을 쓴다면 이렇게 추가할 수도 있다: ['XIGNCODE', 'xigncode', 'Xigncode']
new_words = ['xigncode']

# 3. tokenizer에 새 토큰 추가
num_added_toks = tokenizer.add_tokens(new_words)

if num_added_toks > 0:
    print(f"Added {num_added_toks} new tokens.")

    # 4. 모델의 토큰 임베딩 크기 조정
    # 새 토큰(들)을 위해 무작위로 초기화된 벡터가 추가된다
    transformer_model.resize_token_embeddings(len(tokenizer))

    print(f"New vocabulary size: {len(tokenizer)}")
else:
    print("Tokens already exist in the vocabulary.")

# tokenizer 테스트
test_sentence = "This is a test of xigncode and XIGNCODE3."
tokenized = tokenizer.tokenize(test_sentence)

print(f"Tokenization test: {tokenized}")
# 기대 출력: ['this', 'is', 'a', 'test', 'of', 'xigncode', 'and', 'xigncode', '3']

from torch.utils.data import DataLoader
from sentence_transformers import InputExample, losses, models

# 1. 학습 예제 만들기
# 새 단어에 맥락을 제공하는 문장이 필요하다.
train_examples = [
    # 맥락 부여: "xigncode"는 "anti-cheat"과 관련이 있다
    InputExample(texts=['xigncode is an anti-cheat solution.', 'This software prevents cheating in online games.']),

    # 또 다른 맥락 부여:
    InputExample(texts=['Many gamers are familiar with xigncode.', 'It is a well-known security program.'])

    # 더 잘 배우려면 예제를 훨씬 많이 추가할 것...
]

# 2. 데이터로더와 손실 함수 설정
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)
train_loss = losses.MultipleNegativesRankingLoss(model=model)

# 3. 모델 fine-tuning
num_epochs = 1
warmup_steps = int(len(train_dataloader) * num_epochs * 0.1)

print("Starting fine-tuning...")
model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=num_epochs,
    warmup_steps=warmup_steps,
    output_path='./my-finetuned-model',
    show_progress_bar=True
)

print("Fine-tuning complete. Model saved to './my-finetuned-model'.")

from langchain_huggingface import HuggingFaceEmbeddings

# 로컬에 fine-tuning된 모델 디렉토리를 가리킨다
model_path = "./my-finetuned-model"

# LangChain이 모델과 새 tokenizer를 함께 로드한다
embeddings = HuggingFaceEmbeddings(
    model_name=model_path,
    model_kwargs={'device': 'cpu'}, # 가능하면 'cuda'
    encode_kwargs={'normalize_embeddings': True}
)

# 이제 이 텍스트가 더 정확한 임베딩을 만든다
text = "Tell me about the xigncode anti-cheat."
query_embedding = embeddings.embed_query(text)

print("Successfully created embedding with the new word!")
print(f"Vector dimension: {len(query_embedding)}")