nahj_rag/3_convert_khaemenei-json_to_faiss.py
2026-02-17 16:52:37 +00:00

89 lines
3.4 KiB
Python

import json
import numpy as np
import faiss
import os
def create_faiss_index_from_json():
# مسیر فایل‌ها
json_file_path = './output-khamenei/embedding_FastText_khamenei.json'
faiss_index_path = './data-faiss/faiss_index_khamenei.index'
metadata_file_path = './data-faiss/faiss_index_khamenei_metadata.json'
# --- 1. بارگذاری داده‌ها از JSON ---
with open(json_file_path, 'r', encoding='utf-8') as f:
data = json.load(f)
# فرض بر این است که هر عنصر شامل فیلدهای زیر است:
# {
# "year": "سال",
# "sentence": "متن جمله",
# "embeddings": [0.12, 0.34, ...]
# }
ids = []
dates = []
sentences = []
titles = []
urls = []
embeddings_list = []
for k, item in data.items():
# print(k)
# print(item)
ids.append(item['id'])
dates.append(item['date'])
sentences.append(item['sentence'])
urls.append(item['url'])
titles.append(item['title'])
embeddings_list.append(item['embeddings'])
embeddings = np.array(embeddings_list).astype('float32') # ابعاد: (n, d)
dimension = embeddings.shape[1]
print(f"Loaded {len(embeddings)} embeddings with dimension {dimension}")
# --- 2. ایجاد ایندکس FAISS برای GPU ---
# اگر فقط CPU دارید، از faiss.IndexFlatL2 استفاده کنید.
# اگر GPU دارید، ابتدا ایندکس را روی CPU ایجاد و سپس به GPU انتقال دهید.
cpu_index = faiss.IndexFlatL2(dimension) # معیار فاصله L2 (Euclidean)
# انتقال ایندکس به GPU
if faiss.get_num_gpus() > 0:
print("Using GPU for FAISS index...")
res = faiss.StandardGpuResources()
gpu_index = faiss.index_cpu_to_gpu(res, 0, cpu_index)
else:
print("GPU not available, using CPU.")
gpu_index = cpu_index
# --- 3. افزودن داده‌ها به ایندکس ---
gpu_index.add(embeddings)
print(f"Total vectors indexed: {gpu_index.ntotal}")
# --- 4. ذخیره ایندکس به فایل ---
# برای ذخیره باید به CPU منتقل شود
final_index = faiss.index_gpu_to_cpu(gpu_index) if isinstance(gpu_index, faiss.Index) and faiss.get_num_gpus() > 0 else gpu_index
os.makedirs(os.path.dirname(faiss_index_path), exist_ok=True)
faiss.write_index(final_index, faiss_index_path)
print(f"FAISS index saved to {faiss_index_path}")
# --- 5. ذخیره متادیتا (برای نگاشت نتایج جستجو) ---
metadata = [{
"id": i,
"date": d,
"title": t,
"sentence": s,
"url": u,
} for i, d, t, s, u in zip(ids, dates, titles, sentences, urls)]
# metadata = [{ "sentence": s} for s in zip(titles, sentences)]
with open(metadata_file_path, 'w', encoding='utf-8') as f:
json.dump(metadata, f, ensure_ascii=False, indent=2)
print(f"Metadata saved to {metadata_file_path}")
if __name__ == '__main__':
# استفاده از متد
# json_file_path = './output-speechs/embedding_FastText_khamenei.json' # مسیر فایل JSON
# index_file_path = './data-faiss/faiss_index_khamenei.index' # مسیر فایل ایندکس FAISS
# index_file_path = './data-faiss/faiss_index_khamenei_metadata.json' # مسیر فایل متادیتا FAISS
create_faiss_index_from_json()