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

kimseongsan/ko-sbert-384-v2

Korean SBERT 384-dim

성능 지표

데이터 출처

사용법

학습 세부사항

다음 단계

kimseongsan/
ko-sbert-384-v2