Spaces:
Sleeping
Sleeping
File size: 2,460 Bytes
9eba547 | 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 | import argparse
import os
import sys
import json
import logging
from tqdm import tqdm
# 添加项目根目录到Python路径
root_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(root_dir)
# 然后使用相对于项目根目录的导入
from modules.title_generator.title_generator import RagTitleGenerator
# 配置日志
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("BuildTitleIndex")
def process(data,
index_path: str='faiss_infographics.index',
data_path: str='infographics_data.npy',
embed_model_path: str='',
force: bool=False):
try:
# 检查索引文件是否存在
if os.path.exists(index_path) and not force:
logger.info(f"索引文件 {index_path} 已存在，跳过创建。使用 --force 参数强制重建。")
return True
logger.info("Initializing RagTitleGenerator...")
generator = RagTitleGenerator(
index_path=index_path,
data_path=data_path,
embed_model_path=embed_model_path
)
logger.info("Building FAISS index...")
generator.build_faiss_index(data)
logger.info("Index built and saved successfully.")
return True
except Exception as e:
logger.error(f"构建索引失败: {str(e)}")
return False
def main(force: bool=False):
parser = argparse.ArgumentParser(description="Build FAISS index for title generation")
parser.add_argument('--data', type=str, required=True, help='Path to training data JSON file')
parser.add_argument('--index_path', type=str, default='faiss_infographics.index', help='Path to store FAISS index')
parser.add_argument('--data_path', type=str, default='infographics_data.npy', help='Path to store embedding + title data')
parser.add_argument('--embed_model_path', type=str, default='', help='Path to sentence embedding model (optional)')
parser.add_argument('--force', action='store_true', help='Force rebuild even if index exists')
args = parser.parse_args()
process(args.data, args.index_path, args.data_path, args.embed_model_path, force or args.force)
if __name__ == "__main__":
try:
main()
except Exception as e:
logger.error(f"程序执行失败: {str(e)}")
sys.exit(1)
|