한국어 문서 검색에는 어떤 임베딩을 쓸까? Qwen3·BGE·EmbeddingGemma 2 비교

[
  {"id":"p-001","title":"여름휴가 자격","text":"입사 6개월 이상인 직원은 여름휴가를 신청할 수 있습니다."},
  {"id":"p-002","title":"신청 기한","text":"사용 희망일 14일 전까지 신청해야 합니다."}
]

[
  {"id":"q-001","text":"입사 6개월 차도 여름휴가를 신청할 수 있나요?","gold_ids":["p-001"]},
  {"id":"q-002","text":"퇴직 후 미사용 휴가를 현금으로 받나요?","gold_ids":[]}
]

python -m venv .venv
source .venv/bin/activate
python -m pip install -U pip
python -m pip install 'sentence-transformers>=5.0' numpy FlagEmbedding
python qwen_retrieval_baseline.py --model qwen
python qwen_retrieval_baseline.py --model bge
python qwen_retrieval_baseline.py --model gemma

import argparse, json, time
from pathlib import Path
import numpy as np
MODELS = {"qwen":"Qwen/Qwen3-Embedding-0.6B", "bge":"BAAI/bge-m3", "gemma":"google/embeddinggemma-2"}
INSTRUCTION = "Instruct: Given a Korean policy question, retrieve passages that answer the question\nQuery: "
p = argparse.ArgumentParser(); p.add_argument("--model", choices=MODELS, required=True); args=p.parse_args()
for name in ("corpus.json", "questions.json"):
    if not Path(name).is_file(): raise SystemExit(f"Missing {name}; save the JSON example first.")
docs_data=json.loads(Path("corpus.json").read_text(encoding="utf-8")); questions=json.loads(Path("questions.json").read_text(encoding="utf-8"))
if not docs_data or not questions: raise SystemExit("Both JSON files need records.")
start=time.perf_counter()
if args.model=="bge":
    from FlagEmbedding import BGEM3FlagModel
    model=BGEM3FlagModel(MODELS[args.model], use_fp16=False, devices=["cpu"])
else:
    from sentence_transformers import SentenceTransformer
    kw={"device":"cpu"}
    if args.model=="gemma": kw["config_kwargs"]={"vision_config":None,"audio_config":None}
    model=SentenceTransformer(MODELS[args.model], **kw)
load_s=time.perf_counter()-start
docs=[f"title: {d.get('title','none')} | text: {d['text']}" for d in docs_data]
start=time.perf_counter()
if args.model=="bge": vectors=model.encode(docs,batch_size=2,max_length=8192)["dense_vecs"]
else: vectors=model.encode(docs,batch_size=2,normalize_embeddings=True,convert_to_numpy=True)
vectors=np.asarray(vectors,dtype=np.float32); index_s=time.perf_counter()-start
out=Path(f"{args.model}-vectors.npy"); np.save(out,vectors)
hits=answerable=0; no_answer=[]; rows=[]; start=time.perf_counter()
for q in questions:
    gold=q.get("gold_ids",[])
    if not gold: no_answer.append(q["id"]); continue
    if len(gold)!=1: raise ValueError(f"Exactly one gold ID required: {q['id']}")
    text=q["text"]
    if args.model=="qwen": text=INSTRUCTION+text
    if args.model=="gemma": text="task: search result | query: "+text
    if args.model=="bge": qv=model.encode([text],batch_size=1,max_length=8192)["dense_vecs"][0]
    else: qv=model.encode([text],normalize_embeddings=True,convert_to_numpy=True)[0]
    top=np.argsort(vectors@np.asarray(qv,dtype=np.float32))[-5:][::-1]; ids=[docs_data[i]["id"] for i in top]
    hit=gold[0] in ids; hits+=hit; answerable+=1; rows.append({"query_id":q["id"],"hit_at_5":bool(hit),"top5_ids":ids})
search_s=time.perf_counter()-start
print(json.dumps({"model":args.model,"model_load_s":round(load_s,3),"document_embedding_s":round(index_s,3),"query_search_s":round(search_s,3),"answerable_queries":answerable,"hit_at_5":hits/answerable if answerable else None,"no_answer_query_ids":no_answer,"raw_vector_array_bytes_in_ram":int(vectors.nbytes),"raw_npy_file_bytes_only":out.stat().st_size,"top5":rows},ensure_ascii=False,indent=2))