async def main():
    entities, description_embedding_store, jina_embedder = process_and_embed()
    
    context_builder = LocalSearchMixedContext(
        entities=entities,
        entity_text_embeddings=description_embedding_store,
        embedding_vectorstore_key=EntityVectorStoreKey.ID,
        text_embedder=jina_embedder,
    )

    local_context_params = {
        "top_k_mapped_entities": 5,
        "include_entity_rank": True,
        "embedding_vectorstore_key": EntityVectorStoreKey.ID,
        "max_tokens": 4000,
    }

    search_engine = LocalSearch(
        llm=DeepSeekClientWrapper(deepseek_client),
        context_builder=context_builder,
        llm_params={"max_tokens": 1000, "temperature": 0.0},
        context_builder_params=local_context_params,
        response_type="multiple paragraphs",
    )

    initial_summary = await run_rag_query(search_engine, "总结这篇论文的主要观点，包括研究方法、实验流程和主要结果")
    
    if initial_summary:
        dataset = await generate_and_answer_questions(deepseek_client, search_engine, initial_summary)
        
        with open('prompt_optimization_dataset.yaml', 'w', encoding='utf-8') as yaml_file:
            yaml.dump(dataset, yaml_file, allow_unicode=True, sort_keys=False)
        
        print("数据集已保存为 prompt_optimization_dataset.yaml")

if __name__ == "__main__":
    import nest_asyncio
    nest_asyncio.apply()
    asyncio.run(main())

def main():

    # 创建文档列表
    documents = [{'content': f"Question: {item['question']}\nAnswer: {item['answer']}"} for item in dataset]

    # 创建检索器
    retriever = SimpleRetriever(documents)
   
    # 加载数据集
    dataset = load_dataset('prompt_optimization_dataset.yaml')

    # 执行查询
    question = "多文档问答任务中，模型是如何具体评估其应用能力的？"
    retrieved_docs = retriever.retrieve(question, top_k=3)
    answer = generator.generate(question, retrieved_docs)
    
    print("问题:", question)
    print("答案:", answer)

if __name__ == "__main__":
    main()

import yaml
from openai import OpenAI

# 配置 DeepSeek 语言模型
class DeepSeekLM:
    def __init__(self, api_key, base_url, model):
        self.client = OpenAI(api_key=api_key, base_url=base_url)
        self.model = model
        self.kwargs = {
            "temperature": 0.7,
            "max_tokens": 200,
            "top_p": 1,
            "frequency_penalty": 0,
            "presence_penalty": 0
        }

    def __call__(self, prompt, **kwargs):
        merged_kwargs = {**self.kwargs, **kwargs}
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": "你是一个答案总结助手。请将给定的答案简明扼要地总结为120-180个字符。"},
                {"role": "user", "content": prompt}
            ],
            **merged_kwargs
        )
        return response.choices[0].message.content

# 加载数据集
print("正在加载数据集...")
with open('prompt_optimization_dataset.yaml', 'r', encoding='utf-8') as file:
    dataset = yaml.safe_load(file)
print(f"成功加载数据集，共 {len(dataset)} 条记录")