38 lines
1.0 KiB
Python
38 lines
1.0 KiB
Python
"""ko-sroberta embeddings for template-fit content similarity.
|
|
|
|
Model: jhgan/ko-sroberta-multitask (HuggingFace)
|
|
용도: MDX summary ↔ frame description 코사인 유사도 계산
|
|
32 frame 규모 — 벡터 DB 없이 numpy 만으로 충분
|
|
"""
|
|
import numpy as np
|
|
|
|
_model = None
|
|
|
|
|
|
def _get_model():
|
|
global _model
|
|
if _model is None:
|
|
from sentence_transformers import SentenceTransformer
|
|
_model = SentenceTransformer('jhgan/ko-sroberta-multitask')
|
|
return _model
|
|
|
|
|
|
def embed_texts(texts):
|
|
"""Batch embed. Returns numpy (N, 768)."""
|
|
model = _get_model()
|
|
return model.encode(list(texts), show_progress_bar=False, convert_to_numpy=True)
|
|
|
|
|
|
def cosine(a, b):
|
|
"""Single-vector cosine similarity."""
|
|
an = a / (np.linalg.norm(a) + 1e-12)
|
|
bn = b / (np.linalg.norm(b) + 1e-12)
|
|
return float(np.dot(an, bn))
|
|
|
|
|
|
def embed_map(id_to_text):
|
|
"""{id: text} → {id: vec(768,)}"""
|
|
ids = list(id_to_text.keys())
|
|
vecs = embed_texts([id_to_text[i] for i in ids])
|
|
return {ids[i]: vecs[i] for i in range(len(ids))}
|