Spaces:
Running
Running
File size: 2,292 Bytes
fb05e78 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
import argparse
import json
import logging
from langchain_core.documents import Document
from langchain_text_splitters import CharacterTextSplitter
logger = logging.getLogger(__name__)
def chunk_documents(input_file_path: str, output_file_path: str):
"""
JSONファイルをロードし、ドキュメントをチャンクに分割して、新しいJSONファイルに保存します。
"""
# テキスト分割を初期化します
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=50)
# 入力JSONファイルからドキュメントを読み込みます
with open(input_file_path, 'r', encoding='utf-8') as f:
documents_data = json.load(f)
chunked_documents = []
for doc_data in documents_data:
page_content = doc_data.get("page_content", "")
metadata = doc_data.get("metadata", {})
# page_contentをチャンクに分割します
splits_texts = text_splitter.split_text(page_content)
# 分割ごとに新しいドキュメントチャンクを作成します
for i, text in enumerate(splits_texts):
chunk_metadata = metadata.copy()
chunk_metadata['chunk_index'] = i
chunked_doc = {
"metadata": chunk_metadata,
"page_content": text
}
chunked_documents.append(chunked_doc)
# チャンク化されたドキュメントを出力JSONファイルに保存します
with open(output_file_path, 'w', encoding='utf-8') as f:
json.dump(chunked_documents, f, ensure_ascii=False, indent=2)
logger.info(f"{len(documents_data)}個のドキュメントを{len(chunked_documents)}個のチャンクに正常に分割しました。")
logger.info(f"出力は {output_file_path} に保存されました")
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="JSONファイルからドキュメントをチャンク化します。")
parser.add_argument("input_file", help="入力JSONファイルのパス (例: documents.json)。")
parser.add_argument("output_file", help="出力JSONファイルのパス (例: chunks.json)。")
args = parser.parse_args()
chunk_documents(args.input_file, args.output_file)
|