Python WikipediaのURLから記事内容を抽出しチャンク分割＆テキスト保存する

RAGを試すためのデータが欲しかったので、WikipediaのURLを指定し、コンテンツをチャンク化してtxtに保存するコードを作成。

必要なライブラリ

以下、必要なライブラリ。

pip install langchain langchain_community beautifulsoup4 tiktoken

ソース

import os
import re
import shutil
from langchain.document_loaders import WebBaseLoader
from langchain.text_splitter import CharacterTextSplitter
from typing import List, Dict
from concurrent.futures import ThreadPoolExecutor

def clean_filename(filename: str) -> str:
    """ファイル名から不要な文字を除去"""
    return re.sub(r'[<>:"/\\|?*]', '', filename)

def clean_output_directory(output_dir: str) -> None:
    """出力ディレクトリを初期化"""
    # ディレクトリが存在する場合は削除して再作成
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
    os.makedirs(output_dir)

def load_and_split_url(url_data: Dict[str, str]) -> List[Dict]:
    """URLからコンテンツを取得し、チャンク分割して保存"""
    try:
        # URLからドキュメントを取得
        docs = WebBaseLoader(url_data["url"]).load()
        
        # テキスト分割設定
        text_splitter = CharacterTextSplitter.from_tiktoken_encoder(
            separator="\n",
            chunk_size=2000,
            chunk_overlap=200,
        )
        
        # ドキュメントを分割
        splits = text_splitter.split_documents(docs)
        
        return [{
            "title": url_data["title"],
            "content": doc.page_content,
            "index": i
        } for i, doc in enumerate(splits)]
    
    except Exception as e:
        print(f"Error processing {url_data['url']}: {str(e)}")
        return []

def save_chunk(chunk_data: Dict, output_dir: str) -> None:
    """チャンクをファイルに保存"""
    clean_title = clean_filename(chunk_data["title"])
    filename = f"{clean_title}_{chunk_data['index']:03d}.txt"
    filepath = os.path.join(output_dir, filename)
    
    with open(filepath, "w", encoding="utf-8") as f:
        f.write(chunk_data["content"])

def main():
    # URL一覧
    urls = [
        {
            "title": "【ファイルのプレフィクス】",
            "url": "【URL】"
        },
        # 必要に応じて他のURLを追加
    ]

    # 出力ディレクトリの初期化
    output_dir = "output"
    print(f"Cleaning output directory: {output_dir}")
    clean_output_directory(output_dir)

    # 出力ディレクトリの作成
    output_dir = "output"
    os.makedirs(output_dir, exist_ok=True)

    # 並列処理でURLからコンテンツを取得とチャンク分割
    with ThreadPoolExecutor() as executor:
        chunks_list = list(executor.map(load_and_split_url, urls))
    
    # フラット化
    all_chunks = [chunk for chunks in chunks_list for chunk in chunks]

    # 並列処理でファイル保存
    with ThreadPoolExecutor() as executor:
        executor.map(
            lambda chunk: save_chunk(chunk, output_dir),
            all_chunks
        )

if __name__ == "__main__":
    main()

実行結果

実行するとoutputディレクトリにチャンク化したテキストが入る。