LlamaIndexでBM25Retrieverを試す

!pip install -U llama-index llama-index-retrievers-bm25

!pip install nltk

from nltk.stem import PorterStemmer
from llama_index.legacy.indices.keyword_table.utils import simple_extract_keywords


def tokenize_remove_stopwords(text: str) -> List[str]:
    # lowercase and stem words
    text = text.lower()
    stemmer = PorterStemmer()
    words = list(simple_extract_keywords(text))
    return [stemmer.stem(word) for word in words]


tokenize_remove_stopwords("母子手帳を受け取りたいのですが、手続きを教えてください。")

!pip install sudachipy sudachidict-core sudachidict-core sudachidict-small sudachidict-full

from typing import List
from sudachipy import Dictionary, SplitMode
import requests

# stopwordsの設定。SlothLibのものを使用して、いくつか追加した
url = "http://svn.sourceforge.jp/svnroot/slothlib/CSharp/Version1/SlothLib/NLP/Filter/StopWord/word/Japanese.txt"
stopwords = []
for line in requests.get(url).text.split("\r\n"):
    if len(line) > 1:
        stopwords.append(line)
stopwords.extend(
    [
        "ください",
        " ",
        "　",
    ]
)

# キーワードだけ抽出したいので雑に品詞でフィルタした
pos_list = ["名詞","動詞","形容詞"]

sudachi_tokenizer = Dictionary(dict="full").create()
mode = SplitMode.C

# トークナイザーの定義
def ja_tokenizer(text: str) -> List[str]:
    tokens = sudachi_tokenizer.tokenize(text, mode)
    token_list = []
    for token in tokens:
        if token.surface() in stopwords:
            continue
        if token.part_of_speech()[0] not in pos_list:
            continue
        token_list.append(token.surface())
    return token_list

ja_tokenizer("母子手帳を受け取りたいのですが、手続きを教えてください。")

!wget https://d.line-scdn.net/stf/linecorp/ja/csr/dataset_.zip
!unzip dataset_.zip

import pandas as pd

df = pd.read_excel("dataset_.xlsx")
df.rename(columns={
    'サンプル 問い合わせ文': 'Question',
    'サンプル 応答文': 'Answer',
    'カテゴリ1': 'Category',
}, inplace=True)
df.drop(columns=["ID", "カテゴリ2","出典","<参考>UMカテゴリタグ","<参考>ＵＭサービスメニュー\n（標準的な行政サービス名称）"], inplace=True)
df.rename(columns={"サンプルID":"ID"}, inplace=True)
df

from llama_index.core.schema import TextNode, NodeRelationship, RelatedNodeInfo

nodes = []
for idx, row in df.iterrows():
    id = row["ID"]
    text = row["Answer"].replace("\n\n","\n")
    node = TextNode(text=text, id_=id)
    nodes.append(node)

from llama_index.retrievers.bm25 import BM25Retriever

bm25_retriever = BM25Retriever(nodes=nodes, similarity_top_k=5, tokenizer=ja_tokenizer)

results = bm25_retriever.retrieve("母子手帳を受け取りたいのですが、手続きを教えてください。")

for r in results:
    print("Score:",r.get_score())
    print("ID:", r.id_)
    print("Contents:", r.get_content()[:50])
    print("******")

results = bm25_retriever.retrieve("母子手帳 受け取り 手続き 教えて")

for r in results:
    print("Score:",r.get_score())
    print("ID:", r.id_)
    print("Contents:", r.get_content()[:50])
    print("******")

from llama_index.core import VectorStoreIndex
from llama_index.core.node_parser import SentenceSplitter
from llama_index.retrievers.bm25 import BM25Retriever

splitter = SentenceSplitter(chunk_size=1024)
index = VectorStoreIndex.from_documents(
    documents, transformations=[splitter], embed_model=embed_model
)

vector_retriever = index.as_retriever(similarity_top_k=2)

bm25_retriever = BM25Retriever.from_defaults(
    docstore=index.docstore, similarity_top_k=2
)

!pip install -U llama-index "llama-index-retrievers-bm25==0.1.5"