{"repo": "chaoql/ccf-aiops-code", "files": {"main.py": {"850f4d0ef4d7f35a19489eea337ee2c48ceeed09": {"t": 1785674171.6504545, "n": 3774, "h": "0adb2b32325ea932", "c": "import asyncio\nimport time\nfrom dotenv import load_dotenv, find_dotenv\nfrom dotenv import dotenv_values\nfrom llama_index.core import Settings\nfrom llama_index.embeddings.huggingface import HuggingFaceEmbedding\nfrom llama_index.legacy.llms import OpenAILike as OpenAI\nfrom qdrant_client import models\nfrom tqdm.asyncio import tqdm\nfrom llama_index.embeddings.openai import OpenAIEmbedding\nfrom pipeline.ingestion import build_pipeline, build_vector_store, read_data\nfrom pipeline.qa import read_jsonl, save_answers, save_now, judge_json, save_context_now\nfrom pipeline.rag import QdrantRetriever, generation_with_knowledge_retrieval\nimport os\n\n\ndef list_directories(path):\n    # 获取指定路径下的所有文件和文件夹\n    items = os.listdir(path)\n    # 过滤出文件夹\n    directories = [item for item in items if os.path.isdir(os.path.join(path, item))]\n    return directories\n\n\nasync def main():\n    _ = load_dotenv(find_dotenv())  # 导入环境\n    config = dotenv_values(\".env\")\n\n    path = './data_imageSummary'  # 示例路径（请替换为你的目录路径）\n    folders = list_directories(path)  # 获取文件夹名称列表\n    print(folders)\n\n    # 初始化 LLM 嵌入模型 和 Reranker\n    llm = OpenAI(\n        api_key=config[\"GLM_KEY\"],\n        model=\"glm-4\",\n        api_base=\"https://open.bigmodel.cn/api/paas/v4/\",\n        is_chat_model=True,\n    )\n    # embeding_name = \"gte-Qwen1.5-7B-instruct\"\n    embeding = OpenAIEmbedding(model=\"text-embedding-3-large\")\n    # embeding = HuggingFaceEmbedding(\n    #     model_name=\"BAAI/bge-large-zh-v1.5\",\n    #     cache_folder=\"./BAAI/\",\n    #     embed_batch_size=128,\n    # )\n    Settings.embed_model = embeding\n\n    # 初始化 数据ingestion pipeline 和 vector store\n    client, vector_stores = await build_vector_store(config, folders, reindex=False)\n\n    for folder in folders:\n        collection_info = await client.get_collection(folder)\n        if collection_info.points_count == 0:\n            data = read_data(\"./data_imageSummary/\" + folder)\n            pipeline = build_pipeline(llm, embeding, vector_store=vector_stores[folder])\n            # 暂时停止实时索引\n            await client.update_collection(\n                collection_name=folder,\n                optimizer_config=models.OptimizersConfigDiff(indexing_threshold=0),\n            )\n            await pipeline.arun(documents=data, show_progress=True, num_workers=1)\n            # 恢复实时索引\n            await client.update_collection(\n                collection_name=folder,\n                optimizer_config=models.OptimizersConfigDiff(indexing_threshold=20000),\n            )\n            print(folder + \"：\" + str(len(data)))\n\n    # queries = read_jsonl(\"questionAdd.jsonl\")\n    queries = read_jsonl(\"questionPlus.jsonl\")\n\n    # 生成答案\n    print(\"Start generating answers...\")\n\n    results_name = \"submit_result_ReImageSummary.jsonl\"\n    # results_name = \"context_add.jsonl\"\n    resNum = judge_json(results_name)\n    for index, query in tqdm(enumerate(queries), total=len(queries)):\n        if index + 1 <= resNum:\n            continue\n        result = await generation_with_knowledge_retrieval(query_str=query[\"query\"], llm=llm,\n                                                           document=query[\"document\"], config=config,\n                                                           abbreviation=query[\"abbreviation\"],\n                                                           vector_stores=vector_stores,\n                                                           embeding=embeding,\n                                                           )\n        save_now(query, result, results_name)  # 提前保存\n        # save_context_now(query, result, results_name)\n\nif __name__ == \"__main__\":\n    start_time = time.time()\n    asyncio.run(main())\n    end_time = time.time()\n    execution_time = end_time - start_time\n    print(\"程序运行时间为：\", execution_time / 60.0, \"分钟\")\n"}}}}