1 line
4.2 KiB
JSON
1 line
4.2 KiB
JSON
{"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"}}}} |