AI_Menu_Search / core /embedder.py
Juhaha
HF Spaces 데모 배포 (Streamlit + Qdrant 임베디드, 색인 빌드타임 생성)
fbd1091
Raw
History Blame Contribute Delete
1.59 kB
from sentence_transformers import SentenceTransformer
from config import EMBEDDING_MODEL_NAME
class MenuEmbedder:
"""
한국어 특화 임베딩 모델 래퍼.
첫 실행 시 모델 자동 다운로드 (~500MB), 이후 로컬 캐시 사용.
"""
def __init__(self):
print(f"임베딩 모델 로드 중: {EMBEDDING_MODEL_NAME}")
self.model = SentenceTransformer(EMBEDDING_MODEL_NAME)
print("모델 로드 완료")
def embed_text(self, text: str) -> list:
"""단일 텍스트 임베딩"""
vector = self.model.encode(text, normalize_embeddings=True)
return vector.tolist()
def embed_batch(self, texts: list, batch_size: int = 32) -> list:
"""배치 임베딩"""
vectors = self.model.encode(
texts,
batch_size=batch_size,
normalize_embeddings=True,
show_progress_bar=True
)
return vectors.tolist()
def build_embedding_text(self, desc: dict) -> str:
"""
Description의 여러 필드를 하나의 임베딩용 텍스트로 결합.
embedding_text 필드를 우선 사용하되, 없으면 주요 필드 조합.
"""
if desc.get("embedding_text"):
return desc["embedding_text"]
parts = [
desc.get("menu_name", ""),
desc.get("function_desc", ""),
" ".join(desc.get("keywords", [])),
" ".join(desc.get("user_verbs", [])),
" ".join(desc.get("sample_queries", [])),
]
return " ".join(filter(None, parts))