Spaces:
Sleeping
Sleeping
File size: 4,673 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 |
"""
埋め込み（embedding）の保存・読み込みユーティリティ
"""
import json
import logging
import os
from datetime import datetime
from pathlib import Path
from typing import List, Dict, Tuple
import numpy as np
from langchain_core.documents import Document
logger = logging.getLogger(__name__)
def save_embeddings(
documents: List[Document],
vectors: List[List[float]],
model_name: str,
embeddings_dir: str = 'data/embeddings'
) -> str:
"""
埋め込みとメタデータをJSONファイルに保存する
Args:
documents: Documentオブジェクトのリスト
vectors: 埋め込みベクトルのリスト
model_name: 使用した埋め込みモデル名
embeddings_dir: 埋め込みを保存するディレクトリ
Returns:
保存した埋め込みファイルのパス
"""
# ディレクトリがなければ作成
Path(embeddings_dir).mkdir(parents=True, exist_ok=True)
# タイムスタンプ生成（ファイル名用）
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
# 埋め込みデータの準備
embeddings_data = {
"model": model_name,
"timestamp": datetime.now().isoformat(),
"total_documents": len(documents),
"embeddings": []
}
# Documentとベクトルを結合
for i, (doc, vector) in enumerate(zip(documents, vectors)):
chunk_id = f"doc_{i}_chunk_{doc.metadata.get('chunk_index', 0)}"
embeddings_data["embeddings"].append({
"chunk_id": chunk_id,
"metadata": doc.metadata,
"content": doc.page_content,
"vector": vector # リストとして保存（JSONシリアライズ可能）
})
# JSONファイルに保存
output_file = os.path.join(embeddings_dir, f"embeddings_{timestamp}.json")
with open(output_file, 'w', encoding='utf-8') as f:
json.dump(embeddings_data, f, ensure_ascii=False, indent=2)
# 最新の埋め込みへのシンボリックリンクを作成または更新
latest_link = os.path.join(embeddings_dir, "latest.json")
if os.path.exists(latest_link):
os.remove(latest_link)
try:
os.symlink(os.path.basename(output_file), latest_link)
except OSError:
# シンボリックリンク作成に失敗した場合（例: Windows）、パスをテキストファイルに保存
with open(os.path.join(embeddings_dir, "latest.txt"), 'w') as f:
f.write(output_file)
logger.info(f"Embeddings saved to: {output_file}")
logger.info(f"Total embeddings saved: {len(vectors)}")
return output_file
def load_embeddings(
embeddings_file: str = None,
embeddings_dir: str = 'data/embeddings'
) -> Tuple[List[Document], np.ndarray, str]:
"""
保存済みのJSONファイルから埋め込みを読み込む
Args:
embeddings_file: 読み込む埋め込みファイルのパス。Noneの場合は最新を読み込む
embeddings_dir: 埋め込みが保存されているディレクトリ
Returns:
(documents, ベクトル(numpy配列), モデル名) のタプル
"""
# ファイル指定がなければ最新を読み込む
if embeddings_file is None:
latest_link = os.path.join(embeddings_dir, "latest.json")
if os.path.exists(latest_link):
embeddings_file = latest_link
else:
# latest.txtからパスを取得
latest_txt = os.path.join(embeddings_dir, "latest.txt")
if os.path.exists(latest_txt):
with open(latest_txt, 'r') as f:
embeddings_file = f.read().strip()
else:
raise FileNotFoundError(f"{embeddings_dir} に埋め込みファイルが見つかりません")
# JSONから埋め込みを読み込む
with open(embeddings_file, 'r', encoding='utf-8') as f:
embeddings_data = json.load(f)
# Documentとベクトルを復元
documents = []
vectors = []
for item in embeddings_data["embeddings"]:
# Documentオブジェクトを作成
doc = Document(
page_content=item["content"],
metadata=item["metadata"]
)
documents.append(doc)
vectors.append(item["vector"])
# ベクトルをnumpy配列に変換
vectors_np = np.array(vectors)
logger.info(f"Loaded {len(documents)} embeddings from: {embeddings_file}")
logger.info(f"Model used: {embeddings_data['model']}")
logger.info(f"Created at: {embeddings_data['timestamp']}")
return documents, vectors_np, embeddings_data["model"]
|