bm25

1 post with the tag “bm25”

파이썬으로 BM25와 BM25F 구현하기

Nov 3, 2017

이 포스트에서는 정보 검색 (Information Retrieval) 분야에서 빼 놓을 수 없는 BM25, BM25Simple, BM25F를 파이썬으로 구현해 봅니다. 검색 엔진을 구현해보고자 하는 것이 아니라 BM25, BM25FSimple, BM25F에 대해 코드레벨에서 이해해 보고자 하는 것입니다. 이론은 Okapi BM25에서 확인하세요.

먼저 문서를 준비합니다. 여기서는 몇몇 뉴스기사를 직접 긁어왔습니다. 글이 너무 길어지기 때문에 별도의 파일로 참조합니다.

몇몇 뉴스 기사로 정의한 문서

BM25F의 경우 섹션을 나눠줘야 하는데, 여기서는 간단하게 title, body로 두 개의 섹션을 고려합니다. 다음 코드는 정규식으로 특수문자를 제거하고 BM25를 위한 문서와 BM25F를 적용하기 위한 문서 데이터를 준비합니다.

import re
docs_bm25 = [re.sub('([^a-zA-Z0-9ㄱ-ㅣ가-힣\s])+',' ', doc).split() for doc in docs_bm25_raw]
docs_bm25f = [re.sub('([^a-zA-Z0-9ㄱ-ㅣ가-힣\s])+',' ', doc).split('\n') for doc in docs_bm25_raw]
docs_bm25f = [[title.split(), body.split()] for title, body in docs_bm25f]

다음 코드는 BM25 클래스 입니다.

from functools import reduce
from collections import defaultdict
import numpy as np
import pandas as pd

class BM25:
    def __init__(self, docs):
        self.docs = docs  # array of term vectors
        self.N = 10000  # len(self.docs)
        self.avdl = reduce(lambda dl1, dl2: dl1 + dl2, map(lambda d: len(d), self.docs)) / len(self.docs)

    def rank(self, q, b=0.5, k=2):
        # initialize
        B_d = []  # B constant for docs
        tf_q = [[] for _ in range(len(q))]  # term frequency vector of docs for query terms
        n_q = [0] * len(q)  # doc frequency vectors for query terms

        for idx_d, d in enumerate(self.docs):
            # provision
            d = np.array(d)

            # document length and B
            dl = len(d)
            B_d.append((1-b) / (b * (dl / self.avdl)))

            # term frequency
            for idx_q, queryterm in enumerate(q):
                d_query_included = np.core.defchararray.find(d, queryterm)
                tf_queryterm = len(d_query_included[d_query_included != -1])
                tf_q[idx_q].append(tf_queryterm)
                # updating n
                if tf_queryterm > 0:
                    n_q[idx_q] += 1

        # tf normalization
        tf_q = np.array(tf_q)
        B_d = np.array(B_d)
        tf_q_norm = tf_q / B_d.reshape((1,len(B_d))).repeat(len(q), axis=0)

        # inverse document frequency
        n_q = np.array(n_q)
        idf_q = np.log((self.N - n_q + .5) / (n_q + .5))

        # weight
        w_q = (tf_q_norm / (k + tf_q_norm)) * idf_q.reshape((len(q), 1)).repeat(len(self.docs), axis=1)
        w = w_q.sum(axis=0)

        docidx_ranked = w.argsort(axis=0)

        # result as data frame
        df_result = pd.DataFrame(columns=['docid', 'w', 'doc'], data={
            'docid': list(range(1, len(self.docs)+1)),
            'w': w,
            'doc': self.docs
        })
        return df_result.sort_values(by='w')[::-1].reset_index(drop=True)

다음은 BM25FSimple 클래스 코드 입니다.

from functools import reduce
from collections import defaultdict
import numpy as np
import pandas as pd

class BM25FSimple:
    def __init__(self, docs, v):
        self.docs = docs  # array of field term vectors
        self.v = v  # stream weight
        self.N = 10000  # len(self.docs)
        # calculate avdl after flattening self.docs
        self.avdl = reduce(lambda dl1, dl2: dl1 + dl2, map(lambda d: len(d), np.array(self.docs).flatten())) / len(self.docs)

    def rank(self, q, b=0.5, k=2):
        # initialize
        B_d = []  # B constant for docs
        tf_q = [[] for _ in range(len(q))]  # term frequency vector of docs for query terms
        n_q = [0] * len(q)  # doc frequency vectors for query terms

        for idx_d, d in enumerate(self.docs):
            dl = 0
            tf_queryterm = 0
            for stream, stream_weight in zip(d, self.v):
                # provision
                stream = np.array(stream)

                # document length
                dl += len(stream) * stream_weight

            # term frequency
            for idx_q, queryterm in enumerate(q):
                tf_queryterm = 0
                for stream, stream_weight in zip(d, self.v):
                    stream_query_included = np.core.defchararray.find(stream, queryterm)
                    tf_queryterm += len(stream_query_included[stream_query_included != -1]) * stream_weight
                tf_q[idx_q].append(tf_queryterm)
                # updating n
                if tf_queryterm > 0:
                    n_q[idx_q] += 1

            # B; sum of length of streams
            B_d.append((1-b) / (b * (dl / self.avdl)))

        # tf normalization
        tf_q = np.array(tf_q)
        B_d = np.array(B_d)
        tf_q_norm = tf_q / B_d.reshape((1,len(B_d)))

        # inverse document frequency
        n_q = np.array(n_q)
        idf_q = np.log((self.N - n_q + .5) / (n_q + .5))

        # weight
        w_q = (tf_q_norm / (k + tf_q_norm)) * idf_q.reshape((len(q), 1))
        w = w_q.sum(axis=0)

        docidx_ranked = w.argsort(axis=0)

        # result as data frame
        df_result = pd.DataFrame(columns=['docid', 'w', 'doc'], data={
            'docid': list(range(1, len(self.docs)+1)),
            'w': w,
            'doc': self.docs
        })
        return df_result.sort_values(by='w')[::-1].reset_index(drop=True)

다음 코드는 BM25F 클래스 코드 입니다.

from functools import reduce
from collections import defaultdict
import numpy as np
import pandas as pd

class BM25F:
    def __init__(self, docs, v):
        if len(docs) == 0:
            print("Error. Empty document list.")
            return

        self.docs = docs  # array of field term vectors
        self.v = v  # stream weight
        self.N = 10000  # len(self.docs)
        arr_docs = np.array(self.docs)
        self.avsl = []
        for stream_idx in range(arr_docs.shape[1]):
            self.avsl.append(reduce(lambda sl1, sl2: sl1 + sl2, map(lambda s: len(s), arr_docs[:, stream_idx])) / len(self.docs))

    def rank(self, q, b=.5, k=2):
        if type(b) is not list:
            b = [b] * len(self.docs[0])

        # initialize
        B_d = []  # B constant for docs
        tf_q = [[] for _ in range(len(q))]  # term frequency vector of docs for query terms
        n_q = [0] * len(q)  # doc frequency vectors for query terms

        for idx_d, d in enumerate(self.docs):
            B_s = []
            tf_queryterm = 0
            tf_stream = []

            # initialize tf
            for idx_q in range(len(q)):
                tf_q[idx_q].append(0)

            for stream, stream_weight, avsl, b_s in zip(d, self.v, self.avsl, b):
                # provision
                stream = np.array(stream)

                # document length
                sl = len(stream)

                # B
                B = (1-b_s) / (b_s * (sl / avsl))

                # term frequency per stream
                for idx_q, queryterm in enumerate(q):
                    stream_query_included = np.core.defchararray.find(stream, queryterm)
                    tf_queryterm_stream = len(stream_query_included[stream_query_included != -1]) * stream_weight / B
                    tf_q[idx_q][-1] += tf_queryterm_stream  # add tf of the stream

            # update document frequency
            for idx_q, queryterm in enumerate(q):
                if tf_q[idx_q][-1] > 0:
                    n_q[idx_q] += 1

        # tf normalization
        tf_q = np.array(tf_q)

        # inverse document frequency
        n_q = np.array(n_q)
        idf_q = np.log((self.N - n_q + .5) / (n_q + .5))

        # weight
        w_q = (tf_q / (k + tf_q)) * idf_q.reshape((len(q), 1))
        w = w_q.sum(axis=0)

        docidx_ranked = w.argsort(axis=0)

        # result as data frame
        df_result = pd.DataFrame(columns=['docid', 'w', 'doc'], data={
            'docid': list(range(1, len(self.docs)+1)),
            'w': w,
            'doc': self.docs
        })
        return df_result.sort_values(by='w')[::-1].reset_index(drop=True)

이제 몇몇 질의에 대한 랭킹을 해봅니다.

bm25 = BM25(docs_bm25)
bm25f_simple = BM25FSimple(docs_bm25f, [3, 1])
bm25f = BM25F(docs_bm25f, [3, 1])

q_list = [
    ['네이버'],
    ['네이버', '웹툰'],
    ['네이버', '후원'],
    ['네이버', '웹툰', '후원'],
    ['대통령'],
    ['총리'],
    ['대통령', '총리'],
    ['네이버', '문재인', '대통령']
]

df_result_list = []
for q in q_list:
    df_result_list.append([bm25.rank(q), bm25f_simple.rank(q), bm25f.rank(q)])

이제 Term Frequency (TF)를 계산해 봅니다.

df_result_combined_list = []
for df_result in df_result_list:
    df_result_combined_list.append(
        pd.DataFrame(columns=['docid', 'w_bm25', 'w_bm25f_simple', 'w_bm25f'], data={
            'docid': df_result[0].sort_values(by='docid').set_index('docid').index.values,
            'w_bm25': df_result[0].sort_values(by='docid').set_index('docid')['w'],
            'w_bm25f_simple': df_result[1].sort_values(by='docid').set_index('docid')['w'],
            'w_bm25f': df_result[2].sort_values(by='docid').set_index('docid')['w']
        })
    )

이제 계산한 BM25, BM25FSimple, BM25F 값과 TF를 가시화 합니다.

dflist_query_tf = []
for q in q_list:
    df_query_tf = pd.DataFrame()
    df_query_tf['docid'] = list(range(len(docs_bm25f)))
    df_query_tf = df_query_tf.set_index('docid')
    for queryterm in q:
        title_tf = []
        body_tf = []
        for title, body in docs_bm25f:
            title = np.array(title)
            body = np.array(body)
            title_query_included = np.core.defchararray.find(title, queryterm)
            body_query_included = np.core.defchararray.find(body, queryterm)
            title_tf.append(len(title_query_included[title_query_included != -1]))
            body_tf.append(len(body_query_included[body_query_included != -1]))
        df_query_tf["title."+queryterm] = title_tf
        df_query_tf["body."+queryterm] = body_tf
    dflist_query_tf.append(df_query_tf)

import matplotlib
import matplotlib.pyplot as plt
import matplotlib.font_manager as fm
font_location = "/Library/Fonts/NanumBarunGothic.otf"
font_name = fm.FontProperties(fname=font_location).get_name()
matplotlib.rc('font', family=font_name)
%matplotlib inline

fig, axes = plt.subplots(nrows=len(q_list), ncols=2, figsize=(15, len(q_list)*3))

def plot(ax, df, title):
    ax = df.sort_values(by='docid').set_index('docid').plot.bar(ax=ax, title=title)

for ax, df, q in zip(axes[:, 0], df_result_combined_list, q_list):
    plot(ax, df, 'q=%s' % (q,))

for ax, df in zip(axes[:, 1], dflist_query_tf):
    df.plot.bar(ax=ax)

fig.tight_layout()

가시화 결과입니다.

Tags:

• python
• bm25
• bm25f