LlamaIndexのMarkdown Node Parserを試してみる

from pathlib import Path
import requests
import re

def replace_heading(match):
    level = len(match.group(1))
    return '#' * level + ' ' + match.group(2).strip()

# Wikipediaからのデータ読み込み
wiki_titles = ["イクイノックス", "ドウデュース"]
for title in wiki_titles:
    response = requests.get(
        "https://ja.wikipedia.org/w/api.php",
        params={
            "action": "query",
            "format": "json",
            "titles": title,
            "prop": "extracts",
            # 'exintro': True,
            "explaintext": True,
        },
    ).json()
    page = next(iter(response["query"]["pages"].values()))
    wiki_text = f"# {title}\n\n"
    wiki_text += page["extract"]

    wiki_text = re.sub(r"(=+)([^=]+)\1", replace_heading, wiki_text)
    wiki_text = re.sub(r"\t+", "", wiki_text)
    wiki_text = re.sub(r"\n{3,}", "\n\n", wiki_text)
    data_path = Path("data")
    if not data_path.exists():
        Path.mkdir(data_path)

    with open(data_path / f"{title}.md", "w") as fp:
        fp.write(wiki_text)

!cat イクイノックス.md

# イクイノックス

イクイノックス（欧字名:Equinox、2019年3月23日 - ）は、日本の競走馬。
キタサンブラック産駒として初のGI制覇を果たし、その後GI・6連勝を達成。2023年にはワールド・ベスト・レースホース・ランキングにおいて日本調教馬として史上2頭目の同ランキング単独1位となり、｢世界一」の競走馬となった。 馬名の意味は「昼と夜の長さがほぼ等しくなる時」。
主な勝ち鞍は2022年・2023年の天皇賞（秋）連覇、2022年の有馬記念、2023年のドバイシーマクラシック、宝塚記念、ジャパンカップ。2022年・2023年のJRA賞年度代表馬、2023年のワールドベストレースホースである。

## 血統・デビュー前
キタサンブラックの初年度産駒である。GIを7勝し、演歌歌手・北島三郎が実質的なオーナーである事からも注目を浴びた父を持ち、母はマーメイドステークスを制覇したシャトーブランシュ。その父は高松宮記念を制覇したキングヘイローである。
2019年3月23日、北海道安平町のノーザンファームで誕生。一口馬主法人シルクホースクラブから総額4000万円（一口8万円×500口）で募集され、ノーザンファーム早来で育成。厩舎長の桑田裕規によると、距離適性と馬体の成長面から、当時の目標としては父の制した菊花賞が据えられていた。その後、美浦トレーニングセンターの木村哲也厩舎に入厩した。

## 戦績

### 2歳（2021年）
(snip)

## 血統表
母シャトーブランシュは2015年のマーメイドステークスの勝ち馬。
父の父ブラックタイドは2005年の中央競馬クラシック三冠ディープインパクトの全兄。
半兄にヴァイスメテオール（父：キングカメハメハ）（ラジオNIKKEI賞）、近親にブランディス（中山大障害、中山グランドジャンプ）がいる。

## 脚注

### 注釈

### 出典

## 外部リンク

競走馬成績と情報 netkeiba、スポーツナビ、JBISサーチ、Racing Post

from llama_index.node_parser import MarkdownElementNodeParser, MarkdownNodeParser
from llama_index.readers.file.flat_reader import FlatReader
from pathlib import Path
import glob
import os

files = glob.glob('data/*.md')

docs = []
for f in files:
    doc = FlatReader().load_data(Path(f))
    docs.extend(doc)

parser = MarkdownNodeParser()

nodes = parser.get_nodes_from_documents(docs)

nodes[0]

TextNode(id_='f53d3107-da1f-45e9-8f29-89de1e76aa0a', embedding=None, metadata={'Header 1': 'ドウデュース', 'filename': 'ドウデュース.md', 'extension': '.md'}, excluded_embed_metadata_keys=[], excluded_llm_metadata_keys=[], relationships={<NodeRelationship.SOURCE: '1'>: RelatedNodeInfo(node_id='6749f8b2-883b-4103-8d1b-3606891af3ba', node_type=<ObjectType.DOCUMENT: '4'>, metadata={'filename': 'ドウデュース.md', 'extension': '.md'}, hash='ae40af2e39fef43d6bcf01ade0d3c5e79203eb00f28bca4b5359a0beaa5d1395'), <NodeRelationship.NEXT: '3'>: RelatedNodeInfo(node_id='1cdf1ac0-646a-461b-a282-d2e041c292f5', node_type=<ObjectType.TEXT: '1'>, metadata={'Header 1': 'ドウデュース', 'Header 2': '戦績'}, hash='19686bdb4d9aa3ba4b9c4f25d82440764bb9e04fe8eee7df008905973052c1e5')}, text='ドウデュース\n\nドウデュース（欧字名:Do Deuce、2019年5月7日 - ）は、日本の競走馬。主な勝ち鞍は2021年の朝日杯フューチュリティステークス、2022年の東京優駿、2023年の有馬記念。\n馬名の意味は「する＋テニス用語（勝利目前の意味）」。2021年のJRA賞最優秀2歳牡馬である。', start_char_idx=2, end_char_idx=151, text_template='{metadata_str}\n\n{content}', metadata_template='{key}: {value}', metadata_seperator='\n')

from llama_index.schema import MetadataMode

for n in nodes:
    print(n.get_content(metadata_mode=MetadataMode.LLM))
    print("\n----------")

Header 1: ドウデュース
filename: ドウデュース.md
extension: .md

ドウデュース

ドウデュース（欧字名:Do Deuce、2019年5月7日 - ）は、日本の競走馬。主な勝ち鞍は2021年の朝日杯フューチュリティステークス、2022年の東京優駿、2023年の有馬記念。
馬名の意味は「する＋テニス用語（勝利目前の意味）」。2021年のJRA賞最優秀2歳牡馬である。

----------
Header 1: ドウデュース
Header 2: 戦績
filename: ドウデュース.md
extension: .md

戦績

----------
Header 1: ドウデュース
Header 2: 戦績
Header 3: デビュー前
filename: ドウデュース.md
extension: .md

デビュー前
2019年5月7日、北海道安平町のノーザンファームで誕生。松島正昭が代表を務める株式会社キーファーズの所有馬となり、ノーザンファーム空港牧場で育成の後、栗東トレーニングセンターの友道康夫厩舎に入厩した。

----------
(snip)

----------
Header 1: イクイノックス
Header 2: 脚注
Header 3: 出典
filename: イクイノックス.md
extension: .md

出典

----------
(snip)

from llama_index import ServiceContext, VectorStoreIndex
from llama_index.embeddings import OpenAIEmbedding
from llama_index.llms import OpenAI

service_context = ServiceContext.from_defaults(
    llm=OpenAI(model="gpt-3.5-turbo"),
    embed_model=OpenAIEmbedding()
)

index = VectorStoreIndex(
    nodes,
    show_progress=True,
    service_context=service_context
)

query_engine = index.as_query_engine(
    similarity_top_k=5,
)

node_doc_ids = index.index_struct.nodes_dict.values()
for node in index.docstore.get_nodes(node_doc_ids):
    print(node.id_, ":", node.text[:20].replace("\n", " "))

f53d3107-da1f-45e9-8f29-89de1e76aa0a : ドウデュース  ドウデュース（欧字名:D
1cdf1ac0-646a-461b-a282-d2e041c292f5 : 戦績
18d57dd9-ec8c-48cb-97f0-04b0f921c011 : デビュー前 2019年5月7日、北海道安
036d4e4b-c1fe-4842-b0db-1f3c72a431b3 : 2歳（2021年） 9月5日に小倉競馬場
096dc1d2-1499-494f-b067-61d15e56c1c7 : 3歳（2022年） 3歳初戦として、弥生
6b36f30c-a037-4e5f-acc2-b193f26dd304 : 4歳（2023年） 2023年2月12日
3847967b-082e-44f0-97f5-3867da5490f6 : 競走成績 以下の内容はJBISサーチ、n
e1e919df-143a-4b7c-8573-1d513b9ab8ae : 血統表 母ダストアンドダイヤモンズはアメ
9d13bab5-e4e4-4f42-b53f-810b0408a44f : 脚注
1886d588-32ff-4400-b3bc-f5cf37b30959 : 注釈
66e287f2-e408-4cd2-90b2-9f79e05e6b8a : 出典
2ba8fe77-7870-4473-bfec-96a4570e9cd0 : 外部リンク  競走馬成績と情報 netk
b6874f01-f274-4e10-9275-288c5dc32cc2 : イクイノックス  イクイノックス（欧字名
ff4062d9-7088-4eda-81c1-5217893cff6c : 血統・デビュー前 キタサンブラックの初年
7157a696-cafe-40a5-a728-ed675bfa2f82 : 戦績
8fa0d61b-c921-41a2-95d0-d1a8cd23b587 : 2歳（2021年） デビュー前に木村の自
53a96f83-56d4-4466-b0d7-4e2393c6f404 : 3歳（2022年）
23ab269a-cfda-4b23-ad86-43a09b67ba23 : 春クラシック競走 予定通りトライアル競走
8e748c89-2831-4078-a61d-a6ac30cc7322 : 天皇賞・秋 次走として天皇賞（秋）に出走
8167c187-b9fe-4b66-ad4b-f13c84994bf8 : 有馬記念 次走として有馬記念に出走すると
3857166c-a7e4-419f-8b0b-142bcf0b387b : 4歳（2023年）
2dad412b-2084-4c6a-b5f0-9f49395d0bbc : ドバイシーマクラシック 1月17日、国際
ea649665-de63-49e7-80a3-42f2c8306dcc : 宝塚記念 次走として宝塚記念に出走すると
8f731f4c-34bb-4d6f-a9ca-9f8ae3c70f55 : 天皇賞・秋連覇 次走として連覇のかかる天
e2f9ea66-6e89-4dd3-91e3-2d75651f46b4 : ジャパンカップ 11月26日、予定通りジ
92300b76-8a4c-4e90-a33f-8ac5190e79ec : 種牡馬時代
f576edee-2a08-47d2-87aa-bbb1487edd2a : 供用 競走馬引退後は、北海道安平町の社台
e3ad3b05-89d8-400f-8fbe-0ad0be120136 : 競走成績 以下の内容は、JBISサーチ、
de03f2e4-e1d6-4f26-8470-478f6642eaf7 : 評価 2023年のジャパンカップでリバテ
a19b51dc-56da-412a-907b-76fbfdb3a4bf : 血統表 母シャトーブランシュは2015年
efc9ff0f-fa82-4a6f-b1e1-d2d5d3cf39fc : 脚注
b80d7713-1384-49f1-ad2a-ea865e45f3db : 注釈
b65c8991-7b3c-4c2e-abf2-acf648f11ee8 : 出典
95c9dc3f-0878-4018-8cad-fcfe673fa4c3 : 外部リンク  競走馬成績と情報 netk

index.delete_nodes(node_ids=["b65c8991-7b3c-4c2e-abf2-acf648f11ee8"])

NotImplementedError: Vector indices currently only support delete_ref_doc, which deletes nodes using the ref_doc_id of ingested documents.

from llama_index.schema import MetadataMode

nodes_for_delete = []

for idx, n in enumerate(nodes):
    # メタデータからセクション情報を取り出す。
    metadatas = []
    header_keys = []
    for m in n.metadata:
        if m.startswith("Header"):
            metadatas.append(n.metadata[m])
            header_keys.append(m)
    if len(metadatas) > 0:
        # セクション情報を新たなメタデータに設定
        n.metadata["section"] = metadata_str = " > ".join(metadatas)
        # 古いセクション情報を削除
        for k in header_keys:
            if k.startswith("Header"):
               del n.metadata[k]

    # コンテンツ整形
    contents = n.get_content().split("\n")
    if len(contents) == 1:
        # コンテンツが1つだけ≒セクションタイトルのみの場合は削除対象
        nodes_for_delete.append(idx)
    else:
        # コンテンツの冒頭にあるセクションタイトル部分、及びそれに続く改行を削除
        content_for_delete = []
        for c_idx, c in enumerate(contents):
            if c in (metadatas):
                content_for_delete.append(c_idx)
            elif c in ["", "\n", None]:
                content_for_delete.append(c_idx)
            else:
                break
        contents = [item for i, item in enumerate(contents) if i not in content_for_delete]

    # 整形したコンテンツでノードを書き換え
    n.set_content("\n".join(contents))

    # LLMやEmbeddingsに含めないメタデータなどを設定
    n.excluded_llm_metadata_keys = ["extension"]
    n.excluded_embed_metadata_keys = ["extension"]
    

new_nodes = [item for i, item in enumerate(nodes) if i not in nodes_for_delete]

for idx, n in enumerate(new_nodes):
    print(n.get_content(metadata_mode=MetadataMode.LLM))
    print("\n----------")

filename: ドウデュース.md
section: ドウデュース

ドウデュース（欧字名:Do Deuce、2019年5月7日 - ）は、日本の競走馬。主な勝ち鞍は2021年の朝日杯フューチュリティステークス、2022年の東京優駿、2023年の有馬記念。
馬名の意味は「する＋テニス用語（勝利目前の意味）」。2021年のJRA賞最優秀2歳牡馬である。

----------
filename: ドウデュース.md
section: ドウデュース > 戦績 > デビュー前

2019年5月7日、北海道安平町のノーザンファームで誕生。松島正昭が代表を務める株式会社キーファーズの所有馬となり、ノーザンファーム空港牧場で育成の後、栗東トレーニングセンターの友道康夫厩舎に入厩した。

----------
(snip)
----------
filename: イクイノックス.md
section: イクイノックス > 血統表

母シャトーブランシュは2015年のマーメイドステークスの勝ち馬。
父の父ブラックタイドは2005年の中央競馬クラシック三冠ディープインパクトの全兄。
半兄にヴァイスメテオール（父：キングカメハメハ）（ラジオNIKKEI賞）、近親にブランディス（中山大障害、中山グランドジャンプ）がいる。

----------
filename: イクイノックス.md
section: イクイノックス > 外部リンク

競走馬成績と情報 netkeiba、スポーツナビ、JBISサーチ、Racing Post

from llama_index import ServiceContext, VectorStoreIndex
from llama_index.embeddings import OpenAIEmbedding
from llama_index.llms import OpenAI

service_context = ServiceContext.from_defaults(
    llm=OpenAI(model="gpt-3.5-turbo"),
    embed_model=OpenAIEmbedding()
)

index = VectorStoreIndex(
    new_nodes,
    show_progress=True,
    service_context=service_context
)

query_engine = index.as_query_engine(
    similarity_top_k=5,
)

response = query_engine.query("イクイノックスの主な勝ち鞍について教えて。")
print(response)

イクイノックスの主な勝ち鞍は、ジャパンカップ、天皇賞（秋）、有馬記念の連覇です。

node_doc_ids = index.index_struct.nodes_dict.values()
for node in index.docstore.get_nodes(node_doc_ids):
    print(node.id_, ":", node.text[:30].replace("\n", " "))

6d7c0559-bcf7-4a4d-bb4d-9429bff9ea27 : ドウデュース（欧字名:Do Deuce、2019年5月7日 
a8c35d42-7134-45f4-bc3d-f54b9e76d2af : 2019年5月7日、北海道安平町のノーザンファームで誕生。松
beec9c83-23b0-4576-9063-c81ddf70f135 : 9月5日に小倉競馬場で行われた2歳新馬戦（芝1800メートル
36e5b61a-c70a-4098-affc-c1db0b2d8eba : 3歳初戦として、弥生賞ディープインパクト記念に出走。単勝オッ
bc47abfc-1b3c-4090-acec-1098ee5bde1f : 2023年2月12日に京都競馬場の改修により、阪神競馬場（芝
83a1748b-f84a-4b96-b1a2-901a6ac41116 : 以下の内容はJBISサーチ、netkeiba.com、Fra
b6996a32-ec67-41aa-8afe-2046c6c63e08 : 母ダストアンドダイヤモンズはアメリカで重賞2勝を挙げ、201
82deeeb6-4720-4475-ac47-aecfad2005af : 競走馬成績と情報 netkeiba、スポーツナビ、JBISサ
9776f028-0041-4a37-b9da-9f22edb7545e : イクイノックス（欧字名:Equinox、2019年3月23日
39777a18-e223-4195-93fb-f6f2fab6b3cc : キタサンブラックの初年度産駒である。GIを7勝し、演歌歌手・
58d8d769-f912-4655-92b8-0e9c38cd2373 : デビュー前に木村の自厩舎所属騎手に対するパワーハラスメントの
35886853-6cda-475f-8d04-ac0f559d6574 : 予定通りトライアル競走を用いず、3歳初戦として中147日のロ
9f3e203b-06bc-4cfa-a6d8-74116446aa1d : 次走として天皇賞（秋）に出走することを表明した。東京優駿出走
803ffb1f-752f-409f-91e1-d47b5b10372a : 次走として有馬記念に出走すると表明した。ファン投票でも多くの
f15e0458-177b-4741-819c-f0e0e70de16b : 1月17日、国際競馬統括機関連盟が「ロンジンワールドベストレ
f44109df-9182-4817-b190-ae28d78caea7 : 次走として宝塚記念に出走すると表明した。鞍上は引き続きクリス
a8ecbfcd-5b0b-4850-9c31-55089b292bc5 : 次走として連覇のかかる天皇賞（秋）に出走すると表明した。 1
594b1534-c44a-4d84-8807-a1172bd32307 : 11月26日、予定通りジャパンカップに出走。前半1000mを
e0671301-ec2e-4fa3-a1b6-1124a317f784 : 競走馬引退後は、北海道安平町の社台スタリオンステーションに5
758b0f0b-810f-495a-9a79-bead5c51f83e : 以下の内容は、JBISサーチ、netkeiba.com、エミ
c7ba7ede-b91e-4457-836d-5afd45e19fad : 2023年のジャパンカップでリバティアイランドに騎乗しイクイ
9008f44f-e0b3-4e3b-a4ff-c3d273a5aba1 : 母シャトーブランシュは2015年のマーメイドステークスの勝ち
1010aa44-1a87-49c2-97c4-edd11db9ae7c : 競走馬成績と情報 netkeiba、スポーツナビ、JBISサ

イクイノックスの主な勝ち鞍には、以下のGI（グレード1）競走の勝利が含まれます：

- 天皇賞（秋）：2022年と2023年に連覇を達成し、史上3頭目となる天皇賞（秋）連覇を果たしました。
- ジャパンカップ：2023年に勝利し、GI6連勝を達成。この勝利により、総獲得賞金が22億1544万6100円となり、歴代1位となりました。
- 有馬記念：2022年に勝利し、3歳馬としては前年のエフフォーリア以来2年連続、史上21頭目の有馬記念制覇を成し遂げました。

これらの勝利により、イクイノックスは日本競馬界において非常に高い評価を受けており、多くの記録を打ち立てています。

from pathlib import Path
import requests
import re

def replace_heading(match):
    level = len(match.group(1))
    return '#' * level + ' ' + match.group(2).strip()

# Wikipediaからのデータ読み込み
wiki_titles = ["イクイノックス", "ドウデュース"]
for title in wiki_titles:
    response = requests.get(
        "https://ja.wikipedia.org/w/api.php",
        params={
            "action": "query",
            "format": "json",
            "titles": title,
            "prop": "extracts",
            # 'exintro': True,
            "explaintext": True,
        },
    ).json()
    page = next(iter(response["query"]["pages"].values()))
    wiki_text = f"# {title}\n\n"
    wiki_text += page["extract"]

    wiki_text = re.sub(r"(=+)([^=]+)\1", replace_heading, wiki_text)
    wiki_text = re.sub(r"\t+", "", wiki_text)
    wiki_text = re.sub(r"\n{3,}", "\n\n", wiki_text)
    data_path = Path("data")
    if not data_path.exists():
        Path.mkdir(data_path)

    # markdown(.md)ファイルとして出力
    with open(data_path / f"{title}.md", "w") as fp:
        fp.write(wiki_text)

from llama_index import SimpleDirectoryReader

docs = SimpleDirectoryReader(input_dir="data").load_data()

for d in docs:
    print(d.get_content().replace("\n","")[:100])

イクイノックスイクイノックス（欧字名:Equinox、2019年3月23日 - ）は、日本の競走馬。キタサンブラック産駒として初のGI制覇を果たし、その後GI・6連勝を達成。2023年にはワールド・ベ
-----
血統・デビュー前キタサンブラックの初年度産駒である。GIを7勝し、演歌歌手・北島三郎が実質的なオーナーである事からも注目を浴びた父を持ち、母はマーメイドステークスを制覇したシャトーブランシュ。その父は
-----
戦績
-----
2歳（2021年）デビュー前に木村の自厩舎所属騎手に対するパワーハラスメントの一件で略式命令を受けたため、JRAから調教停止処分を受けた事に伴い、2021年7月29日から10月31日まで岩戸孝樹厩舎に
-----
3歳（2022年）
-----
春クラシック競走予定通りトライアル競走を用いず、3歳初戦として中147日のローテーションで皐月賞に出走。前年のJRA賞最優秀2歳牡馬であるドウデュース、無敗で共同通信杯を制したダノンベルーガに次ぐ3番
(snip)

from llama_index.schema import MetadataMode

nodes_for_delete = []
sections_for_delete = ["競走成績","外部リンク"]

for idx, n in enumerate(nodes):
    # メタデータからセクション情報を取り出す。
    metadatas = []
    header_keys = []
    for m in n.metadata:
        if m.startswith("Header"):
            metadatas.append(n.metadata[m])
            header_keys.append(m)
    if len(metadatas) > 0:
        # セクション情報を新たなメタデータに設定
        n.metadata["section"] = metadata_str = " > ".join(metadatas)
        # 古いセクション情報を削除
        for k in header_keys:
            if k.startswith("Header"):
               del n.metadata[k]

    # コンテンツ整形
    contents = n.get_content().split("\n")
    if len(contents) == 1:
        # コンテンツが1つだけ≒セクションタイトルのみの場合は削除対象
        nodes_for_delete.append(idx)
    elif contents[0] in sections_for_delete:
        # 任意のセクションを削除対象
        nodes_for_delete.append(idx)
    else:
        # コンテンツの冒頭にあるセクションタイトル部分、及びそれに続く改行を削除
        content_for_delete = []
        for c_idx, c in enumerate(contents):
            if c in (metadatas):
                content_for_delete.append(c_idx)
            elif c in ["", "\n", None]:
                content_for_delete.append(c_idx)
            else:
                break
        contents = [item for i, item in enumerate(contents) if i not in content_for_delete]

    # 整形したコンテンツでノードを書き換え
    n.set_content("\n".join(contents))

    # LLMやEmbeddingsに含めないメタデータなどを設定
    n.excluded_llm_metadata_keys = ["extension"]
    n.excluded_embed_metadata_keys = ["extension"]

base_nodes = [item for i, item in enumerate(nodes) if i not in nodes_for_delete]