- WordLlama는 LLM의 토큰 임베딩을 재활용해 퍼지 중복 제거, 유사도 계산, 랭킹, 클러스터링, 의미 기반 텍스트 분할을 수행하는 빠르고 가벼운 NLP 툴킷임
- 추론은 토큰 조회와 평균 풀링 중심으로 동작하며, NumPy만으로 실행 가능한 경량 파이프라인과 CPU 최적화를 내세움
- 기본 모델은 256차원 16MB이며, Matryoshka 표현으로 차원을 줄일 수 있고 바이너리 임베딩은 Hamming 유사도로 더 빠른 계산을 지원함
- MTEB 표에서 WL64~WL1024는 GloVe 300d와 Komninos보다 여러 지표에서 높은 점수를 보이며,
all-MiniLM-L6-v2보다는 대체로 낮은 점수를 기록함
pip install wordllama 후 WordLlama.load()로 사용할 수 있고, .key(query)는 sorted, min, max 같은 표준 라이브러리 함수에 넣을 수 있는 호출 가능 함수를 반환함
WordLlama가 하는 일
- WordLlama는 퍼지 중복 제거, 유사도 계산, 랭킹, 클러스터링, 의미 기반 텍스트 분할 같은 NLP 유틸리티 작업을 위한 경량 툴킷임
- LLaMA 2, LLaMA 3 70B 같은 최신 LLM에서 토큰 임베딩 코드북을 추출해, GloVe·Word2Vec·FastText와 비슷한 compact word representation을 만듦
- 추론 시 의존성이 적고 CPU 하드웨어에 최적화되어 있어 리소스 제약 환경 배포에 적합함
- 빠르고 작은 크기 덕분에 탐색적 분석, LLM 출력 평가기, multi-hop 또는 agentic workflow의 준비 작업 같은 유틸리티 용도로 사용할 수 있음
설치와 기본 사용
pip install wordllama
- 기본 256차원 모델은
WordLlama.load()로 로드함
from wordllama import WordLlama
wl = WordLlama.load()
.key(query)는 Callable[[str], float]를 반환해 후보 문자열을 쿼리와의 유사도 기준으로 정렬하거나 최댓값을 고를 수 있음
query = "Machine learning methods"
candidates = [
"Foundations of neural science",
"Introduction to neural networks",
"Cooking delicious pasta at home",
"Introduction to philosophy: logic",
]
sim_key = wl.key(query)
sorted_candidates = sorted(candidates, key=sim_key, reverse=True)
best_candidate = max(candidates, key=sim_key)
- 예시 결과에서
"Introduction to neural networks"가 점수 0.3414로 가장 높은 후보가 됨
주요 기능
- 임베딩 생성: 단순 토큰 조회와 평균 풀링으로 텍스트 임베딩을 빠르게 생성함
- 유사도 계산: 두 텍스트 사이의 cosine similarity를 계산함
- 문서 랭킹: 쿼리와 후보 문서의 유사도를 기준으로 순위를 매김
- 퍼지 중복 제거: 유사도 임계값을 기준으로 중복 텍스트를 제거함
- 클러스터링: KMeans로 문서를 그룹화함
- 필터링: 쿼리와의 유사도가 기준 이상인 문서만 남김
- Top-K 검색: 쿼리와 가장 유사한 K개 문서를 반환함
- 의미 기반 텍스트 분할: 텍스트를 의미적으로 응집된 청크로 나눔
- 바이너리 임베딩: Hamming 유사도로 더 빠른 계산을 지원함
- Matryoshka 표현: 필요에 따라 임베딩 차원을 잘라 모델 크기와 성능을 조절함
모델 구조와 성능
- WordLlama는 범용 임베딩 프레임워크 안에서 context-less 소형 모델을 학습함
- 기본 모델은 256차원 16MB 크기임
- README의 MTEB 표는 WL64, WL128, WL256, WL512, WL1024를 GloVe 300d, Komninos,
all-MiniLM-L6-v2와 비교함
- WL256은 Clustering
33.25, Reranking 52.03, Classification 58.21, Pair Classification 78.22, STS 67.91, CQA DupStack 24.12, SummEval 30.99를 기록함
- GloVe 300d는 같은 항목에서 각각
27.73, 43.29, 57.29, 70.92, 61.85, 15.47, 28.87을 기록함
all-MiniLM-L6-v2는 Clustering 42.35, Reranking 58.04, Classification 63.05, Pair Classification 82.37, STS 78.90, CQA DupStack 41.32, SummEval 30.81을 기록함
- l2_supercat은 LLaMA 2 vocabulary 모델임
- LLaMA 2 70B와 phi 3 medium 등 여러 모델의 codebook에서 추가 special token을 제거한 뒤 연결해 학습함
- LLaMA 2 tokenizer를 사용한 여러 모델의 codebook을 함께 연결해 학습할 수 있음
- LLaMA 3 70B codebook 학습과 비슷한 성능을 보이면서 vocabulary는
32k 대 128k로 4배 작음
- LLaMA 3 기반 모델로 l3_supercat이 제공됨
- 추가 결과는 Results에 있음
의미 기반 텍스트 분할
.split()은 긴 텍스트를 의미적 청크로 나눔
long_text = "Your very long text goes here... " * 100
chunks = wl.split(long_text, target_size=1536)
print(list(map(len, chunks)))
# Output: [1055, 1055, 1187]
target_size는 목표 크기이면서 최대 크기임- 분할 과정은 텍스트 순서, 문장 구조, 가능한 경우 문단 구조를 유지하려고 함
- WordLlama 임베딩을 사용해 더 자연스러운 분할 인덱스를 찾음
- 출력 청크 크기는
target_size 이하 범위에서 달라질 수 있음
- 권장 target size는 512~2048자이며 기본값은
1536임
- 더 큰 청크가 필요하면 분할 후 여러 semantic chunk를 배치로 묶는 방식을 권장함
- 자세한 내용은 technical overview에 있음
Model2Vec와 직접 추론
- 2025-01-04 업데이트로 Model2Vec static embeddings 지원이 추가됨
WordLlama.load_m2v()로 Model2Vec 모델을 로드할 수 있음
wl = WordLlama.list_configs()
wl = WordLlama.load_m2v("potion_base_8m") # 256-dim model
wl = WordLlama.load_m2v("m2v_multilingual") # multilingual model
- Model2Vec은 PCA를 사용해 static embedding을 만드는 다른 방식임
- Model2Vec 쪽은 multilingual model과 glove 기반 모델을 만들었고, word similarity task에서 좋은 점수를 낸다고 되어 있음
- Hugging Face의 minishlab에서 확인할 수 있음
WordLlamaInference는 로더 대신 (n_vocab, dim) 형태의 static embedding 배열과 tokenizer를 직접 넣어 사용할 수 있음
from wordllama import WordLlamaInference
from tokenizers import Tokenizer
tokenizer = Tokenizer.from_pretrained(...)
wl = WordLlamaInference(np_embeddings_ar, tokenizer)
학습과 임베딩 추출
- 바이너리 임베딩 모델은 고차원에서 개선 폭이 더 뚜렷했고, 바이너리 임베딩에는 512 또는 1024차원이 권장됨
- L2 Supercat 모델은 단일 A100 GPU에서 batch size
512로 12시간 학습됨
- LLaMA 모델에서 토큰 임베딩을 추출하려면 사용자 계약에 동의하고 Hugging Face CLI에 로그인해야 함
from wordllama.extract.extract_safetensors import extract_safetensors
extract_safetensors("llama3_70B", "path/to/saved/model-0001-of-00XX.safetensors")
- 임베딩은 보통 첫 번째
safetensors 파일에 있지만 항상 그렇지는 않음
- manifest가 있을 수도 있음
- 직접 살펴보고 찾아야 할 수도 있음
- 학습에는 저장소의 스크립트를 사용하며, 기존 설정을 복사하거나 수정해 configuration file을 추가해야 함
pip install wordllama[train]
python train.py train --config your_new_config
python train.py save --config your_new_config --checkpoint ... --outdir /path/to/weights/
- 저장 단계는 Matryoshka 차원별로 모델을 하나씩 저장함
업데이트, 로드맵, 라이선스
- 2025-02-01 업데이트로
sorted, min, max 같은 표준 라이브러리 함수에 쓸 수 있는 callable 지원이 추가됨
- 2024-10-04 업데이트로 semantic splitting inference algorithm이 추가됨
- 로드맵에는 DSPy evaluator와 Retrieval-Augmented Generation, 즉 RAG pipeline 예제 노트북 추가가 있음
- 커뮤니티 프로젝트로 Gradio Demo HF Space와 CPU-ish RAG가 있음
- 프로젝트 라이선스는 MIT License임