ko-sbert-384-v2

Korean SBERT 384-dim

jhgan/ko-sbert-multitask의 차원을 768에서 384로 축소한 모델입니다.

성능 지표

• 📐 차원: 768 → 384 (50% 감소)
• 🎯 유사도 보존율: 99.77%
• 📊 최종 손실: 0.000205
• 🔢 학습 샘플: 800개

데이터 출처

• KorNLI (자연어 추론)
• KorSTS (의미 유사도)
• NSMC (영화 리뷰)
• KorQuAD (질의응답)

사용법

import torch
from transformers import AutoModel, AutoTokenizer

class DimensionReducer(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(768, 384)
        self.layer_norm = torch.nn.LayerNorm(384)
    
    def forward(self, x):
        return self.layer_norm(self.linear(x))

# 모델 로드
model = AutoModel.from_pretrained("kimseongsan/ko-sbert-384")
tokenizer = AutoTokenizer.from_pretrained("kimseongsan/ko-sbert-384")

# Reducer 로드
reducer = DimensionReducer()
reducer.load_state_dict(torch.load("reducer.pt"))
reducer.eval()

def encode(sentences):
    if isinstance(sentences, str):
        sentences = [sentences]
    
    inputs = tokenizer(sentences, padding=True, truncation=True, 
                      max_length=128, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model(**inputs)
        attention_mask = inputs['attention_mask']
        token_embeddings = outputs.last_hidden_state
        
        # Mean pooling
        input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
        sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, 1)
        sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)
        embeddings = sum_embeddings / sum_mask
        
        # 차원 축소
        reduced = reducer(embeddings)
    
    return reduced

# 예시
sentences = ["안녕하세요", "반갑습니다"]
embeddings = encode(sentences)
print(f"Shape: {embeddings.shape}")  # torch.Size([2, 384])

학습 세부사항

• Optimizer: Adam (lr=0.001)
• Loss: MSE on similarity matrices
• Epochs: 100
• Batch Size: 32
• Device: cuda:0

다음 단계

이 모델을 INT8 양자화하려면:

python quantize_model.py --model kimseongsan/ko-sbert-384