Spaces:
Sleeping
Sleeping
File size: 4,526 Bytes
f446d43 d72bce3 f446d43 41e01be f446d43 d72bce3 f446d43 41e01be f446d43 d72bce3 f446d43 41e01be f446d43 | 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 136 137 138 139 140 141 142 143 144 145 146 147 | import os
import sys
import json
from typing import Dict, List, Tuple
import numpy as np
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from utils.logger import setup_logger
from utils.json import field_getter, json_dumps
log = setup_logger(__name__)
def load_configs():
files = field_getter("config/files.json")
paths = {
"tf_token": files("bm25.tf_token"),
"fasttext_vec": files("embeddings.fasttext_vec"),
"word_vocab": files("embeddings.word_vocab"),
"word_vectors": files("embeddings.word_vectors"),
}
return paths
def read_vocab_from_tf_token(tf_token_path: str) -> Tuple[List[dict], set[str]]:
with open(tf_token_path, encoding="utf-8") as f:
docs = json.load(f)
vocab: set[str] = set()
for d in docs:
# doc全体tfから語彙を得る
for t in (d.get("tf") or {}).keys():
vocab.add(t)
return docs, vocab
def stream_fasttext_vec(
vec_path: str, vocab: set[str]
) -> Tuple[Dict[str, int], np.ndarray]:
"""fastText .vec から、必要語彙のみ抽出してベクトル行列を返す。"""
token_to_idx: Dict[str, int] = {}
vectors: List[np.ndarray] = []
dim = None
kept = 0
with open(vec_path, encoding="utf-8", errors="ignore") as f:
header = f.readline()
# ヘッダ行は "<count> <dim>" のことが多い
try:
parts = header.strip().split()
if len(parts) >= 2 and parts[0].isdigit():
dim = int(parts[1])
except Exception:
pass
for line in f:
sp = line.rstrip().split(" ")
if len(sp) < 2:
continue
token = sp[0]
if token not in vocab:
continue
vec_vals = sp[1:]
if dim is None:
dim = len(vec_vals)
if len(vec_vals) != dim:
continue
try:
v = np.asarray([float(x) for x in vec_vals], dtype=np.float32)
except ValueError:
continue
# L2正規化
norm = np.linalg.norm(v)
if norm > 0:
v = v / norm
token_to_idx[token] = kept
vectors.append(v)
kept += 1
if dim is None:
raise RuntimeError(".vec の次元を特定できませんでした")
if not vectors:
log.warning("語彙に一致するベクトルが見つかりませんでした")
arr = np.zeros((0, dim), dtype=np.float32)
return token_to_idx, arr
arr = np.vstack(vectors).astype(np.float32)
return token_to_idx, arr
# top-k モードでは文書ベクトルは不要
def main():
log.info(
"語彙/文書ベクトル(word_vocab.json, word_vectors.npz, doc_vectors.npy)を生成します"
)
try:
paths = load_configs()
except Exception as e:
log.error(f"設定の読み込みに失敗しました: {e}")
sys.exit(1)
tf_token_path = paths["tf_token"]
fasttext_vec_path = paths["fasttext_vec"]
if not os.path.exists(fasttext_vec_path):
log.warning(f".vec が見つかりません: {fasttext_vec_path}")
log.warning("Step 6 をスキップします")
sys.exit(0)
try:
docs, vocab = read_vocab_from_tf_token(tf_token_path)
except Exception as e:
log.error(f"tf_token.jsonの読み込みに失敗しました: {e}")
sys.exit(1)
log.info(f"コーパス語彙数: {len(vocab)}")
token_to_idx, word_vecs = stream_fasttext_vec(fasttext_vec_path, vocab)
log.info(f"抽出済み語彙ベクトル数: {word_vecs.shape[0]}")
# 語彙インデックスの安定化（token_to_idxは追加順次第なのでソート）
sorted_tokens = sorted(token_to_idx.keys())
remap = {t: i for i, t in enumerate(sorted_tokens)}
remapped_vecs = np.zeros_like(word_vecs)
for t, old_i in token_to_idx.items():
new_i = remap[t]
remapped_vecs[new_i] = word_vecs[old_i]
token_to_idx = remap
word_vecs = remapped_vecs
# 出力
os.makedirs(os.path.dirname(paths["word_vocab"]), exist_ok=True)
json_dumps(token_to_idx, paths["word_vocab"]) # 語→index
# 圧縮npz
np.savez_compressed(paths["word_vectors"], vectors=word_vecs)
log.info(f"word_vocab.json: {paths['word_vocab']}")
log.info(f"word_vectors.npz: {paths['word_vectors']}")
if __name__ == "__main__":
main()
|