개발/AI

[BM25] 사용 예제

jykim23 2024. 11. 5. 15:27

영어

import nltk
from nltk.tokenize import word_tokenize
from rank_bm25 import BM25Okapi

# NLTK의 punkt 리소스를 다운로드합니다.
nltk.download('punkt')

# 문서 집합을 정의합니다.
documents = [
    "The cat sat on the mat.",
    "Dogs are great pets.",
    "Cats and dogs can be friends.",
    "The quick brown fox jumps over the lazy dog."
]

# 문서 토큰화
tokenized_documents = [word_tokenize(doc.lower()) for doc in documents]

# BM25 모델 초기화
bm25 = BM25Okapi(tokenized_documents)

# 쿼리 정의 및 토큰화
query = "cat pets"
tokenized_query = word_tokenize(query.lower())

# BM25 점수 계산
scores = bm25.get_scores(tokenized_query)

# 결과 출력
for i, score in enumerate(scores):
    print(f"Document {i}: Score = {score}")

 

한국어(openjdk 1.8 설치 필요)

from konlpy.tag import Okt
from rank_bm25 import BM25Okapi

# 한국어 토크나이저 초기화
okt = Okt()

# 문서 집합을 정의합니다.
documents = [
    "고양이가 매트 위에 앉아 있습니다.",
    "개는 훌륭한 애완동물입니다.",
    "고양이와 개는 친구가 될 수 있습니다.",
    "빠른 갈색 여우가 게으른 개를 뛰어넘습니다."
]

# 문서 토큰화
tokenized_documents = [okt.morphs(doc) for doc in documents]

# BM25 모델 초기화
bm25 = BM25Okapi(tokenized_documents)

# 쿼리 정의 및 토큰화
query = "고양이 애완동물"
tokenized_query = okt.morphs(query)

# BM25 점수 계산
scores = bm25.get_scores(tokenized_query)

# 결과 출력
for i, score in enumerate(scores):
    print(f"문서 {i}: 점수 = {score}")