高级RAG技术第1部分：数据处理

# 创建一个名为'rag_env'的新虚拟环境
python -m venv rag_env
# 激活虚拟环境（对于基于Unix的系统）
source rag_env/bin/activate
# （对于Windows）
.\rag_env\Scripts\activate
# 安装requirements.txt中列出的包
pip install -r requirements.txt

# Elastic Cloud: 在Elastic Cloud控制台的“Deployment”页面找到
ELASTIC_CLOUD_ENDPOINT=""
ELASTIC_CLOUD_ID=""
# Elastic Cloud: 在部署设置期间或在“Security”设置中创建
ELASTIC_USERNAME=""
ELASTIC_PASSWORD=""
# Elastic Cloud: 在Kibana或通过API创建的索引名称
ELASTIC_INDEX_NAME=""
# Azure AI Studio: 在Azure OpenAI资源的“Keys and Endpoint”部分找到
AZURE_OPENAI_KEY_1=""
AZURE_OPENAI_KEY_2=""
AZURE_OPENAI_REGION=""
AZURE_OPENAI_ENDPOINT=""
# Azure AI Studio: 在Azure OpenAI资源的“Deployments”部分找到
AZURE_OPENAI_DEPLOYMENT_NAME=""
# 使用BAAI/bge-small-en-v1.5，因为我认为它在资源效率和性能之间取得了良好的平衡。
HUGGINGFACE_EMBEDDING_MODEL="BAAI/bge-small-en-v1.5"

# llamaindex_processor.py
from llama_index.core import SimpleDirectoryReader

class LlamaIndexProcessor:
    def __init__(self):
        pass

    def load_documents(self, directory_path):
        '''加载目录中的所有文档'''
        reader = SimpleDirectoryReader(input_dir=directory_path)
        return reader.load_data()

# main.ipynb
llamaindex_processor = LlamaIndexProcessor()
documents = llamaindex_processor.load_documents('./documents/')
documents = [dict(doc_obj) for doc_obj in documents]

{
  'id_': '5f76f0b3-22d8-49a8-9942-c2bbab14f63f',
  'metadata': {
    'page_label': '5',
    'file_name': 'Elastic_NV_Annual-Report-Fiscal-Year-2023.pdf',
    'file_path': '/Users/han/Desktop/Projects/truckasaurus/documents/Elastic_NV_Annual-Report-Fiscal-Year-2023.pdf',
    'file_type': 'application/pdf',
    'file_size': 3724426,
    'creation_date': '2024-07-27',
    'last_modified_date': '2024-07-27'
  },
  'text': '目录\n页码\n第一部分\n项目1. 业务 3\n15 项目1A. 风险因素\n项目1B. 未解决的员工意见 48\n项目2. 物业 48\n项目3. 法律诉讼 48\n项目4. 矿山安全披露 48\n第二部分\n项目5. 登记人普通股的市场、相关股东事项和发行人股票购买 49\n项目6. [保留] 49\n项目7. 财务状况和经营成果的管理层讨论与分析 50\n项目7A. 关于市场风险的定量和定性披露 64\n项目8. 财务报表和补充数据 66\n项目9. 关于会计和财务披露的会计师变更和分歧 100\n100\n101 项目9A. 控制和程序\n项目9B. 其他信息\n项目9C. 关于防止检查的外国司法管辖区的披露 101\n第三部分\n102\n102\n102\n102 项目10. 董事、高级管理人员和公司治理\n项目11. 高级管理人员薪酬\n项目12. 某些受益所有人和管理层的证券持有情况和相关股东事项\n项目13. 某些关系和相关交易及董事独立性\n项目14. 主要会计师费用和服务 102\n第四部分\n103\n105 项目15. 附件和财务报表附表\n项目16. 10-K表格摘要\n签名 106\ni',
  ...
}

# chunker.py
import uuid
import re

class Chunker:
    def __init__(self, tokenizer):
        self.tokenizer = tokenizer

    def split_into_sentences(self, text):
        """将文本拆分成句子。"""
        return re.split(r'(?<=[.!?])\s+', text)

    def sentence_wise_tokenized_chunk_documents(self, documents, chunk_size=512, overlap=20, min_chunk_size=50):
        '''
        1. 将文本拆分成句子。
        2. 使用提供的分词器方法进行分词。
        3. 构建最多chunk_size限制的块。
        4. 基于令牌创建重叠 - 保留上下文。
        5. 只保留符合最小令牌大小要求的块。
        '''
        chunked_documents = []

        for doc in documents:
            sentences = self.split_into_sentences(doc['text'])
            tokens = []
            sentence_boundaries = [0]

            # 分词所有句子并跟踪句子边界
            for sentence in sentences:
                sentence_tokens = self.tokenizer.encode(sentence, add_special_tokens=True)
                tokens.extend(sentence_tokens)
                sentence_boundaries.append(len(tokens))

            # 创建块
            chunk_start = 0
            while chunk_start < len(tokens):
                chunk_end = chunk_start + chunk_size

                # 找到块中合适的最后一个完整句子
                sentence_end = next((i for i in sentence_boundaries if i > chunk_end), len(tokens))
                chunk_end = min(chunk_end, sentence_end)

                # 创建块
                chunk_tokens = tokens[chunk_start:chunk_end]

                # 检查块是否符合最小大小要求
                if len(chunk_tokens) >= min_chunk_size:
                    # 为此块创建一个新的文档对象
                    chunk_doc = {
                        'id_': str(uuid.uuid4()),
                        'chunk': chunk_tokens,
                        'original_text': self.tokenizer.decode(chunk_tokens),
                        'chunk_index': len(chunked_documents),
                        'parent_id': doc['id_'],
                        'chunk_token_count': len(chunk_tokens)
                    }
                    # 从原始文档复制所有其他字段
                    for key, value in doc.items():
                        if key != 'text' and key not in chunk_doc:
                            chunk_doc[key] = value
                    chunked_documents.append(chunk_doc)

                # 移动到下一个块开始，考虑重叠
                chunk_start = max(chunk_start + chunk_size - overlap, chunk_end - overlap)

        return chunked_documents

# main.ipynb
# 初始化嵌入模型
HUGGINGFACE_EMBEDDING_MODEL = os.environ.get('HUGGINGFACE_EMBEDDING_MODEL')
embedder = EmbeddingModel(model_name=HUGGINGFACE_EMBEDDING_MODEL)

# 初始化分块器
chunker = Chunker(embedder.tokenizer)

chunked_documents = chunker.sentence_wise_tokenized_chunk_documents(documents, chunk_size=512)

print(chunked_documents[4]['original_text'])
[CLS] the aggregate market value of the ordinary shares held by non - affiliates of the registrant, based on the closing price of the shares of ordinary shares on the new york stock exchange on october 31, 2022 ( the last business day of the registrant ’ s second fiscal quarter ), was approximately $ 6. 1 billion. [SEP] [CLS] as of may 31, 2023, the registrant had 97, 390, 886 ordinary shares, par value €0. 01 per share, outstanding. [SEP] [CLS] documents incorporated by reference portions of the registrant ’ s definitive proxy statement relating to the registrant ’ s 2023 annual general meeting of shareholders are incorporated by reference into part iii of this annual ......

# documentenricher.py
from tqdm import tqdm

class DocumentEnricher:
    def __init__(self):
        pass

    def enrich_document(self, documents, processors, text_col='text'):
        for doc in tqdm(documents, desc="Enriching documents using processors: " + str(processors)):
            for (processor, field) in processors:
                metadata = processor(doc[text_col])
                if isinstance(metadata, list):
                    metadata = '\n'.join(metadata)
                doc.update({field: metadata})

# main.ipynb
# 初始化处理器类
nltkprocessor = NLTKProcessor()  # nltk_processor.py
entity_extractor = EntityExtractor()  # entity_extractor.py
gpt4o = LLMProcessor(model='gpt-4o')  # llm.py

# 初始化文档增强器
documentenricher = DocumentEnricher()

# 在文档中创建新的字段 - 这些是处理器函数的输出。
processors = [
    (nltkprocessor.textrank_phrases, "keyphrases"),
    (gpt4o.generate_questions, "potential_questions"),
    (entity_extractor.extract_entities, "entities")
]

# .enrich_document() 将会直接修改 chunked_docs。
# 为了查看结果，我们将在接下来的几个单元中打印 chunked_docs！
documentenricher.enrich_document(chunked_docs, text_col='original_text', processors=processors)

print(chunked_documents[25]['keyphrases'])
'elastic agent stop', 'agent stop malware', 'stop malware ransomware', 'malware ransomware environment', 'ransomware environment wide', 'environment wide visibility', 'wide visibility threat', 'visibility threat detection', 'sep cl key', 'cl key feature'

print(chunked_documents[25]['potential_questions'])
1. Elastic Agent 在网络安全方面的主要功能是什么？
2. 描述 Logstash 如何在 IT 环境中贡献数据管理。
3. 列出并解释文档中提到的 Logstash 的关键特性。
4. Elastic Agent 如何增强威胁检测中的环境可见性？
5. Logstash 提供哪些超越简单数据收集的功能？
6. 文档中如何建议 Elastic Agent 阻止恶意软件和勒索软件？
7. 能否识别 Elastic Agent 和 Logstash 在集成环境中的功能关系？
8. Elastic Agent 的高级威胁检测能力对组织安全政策有何影响？
9. 比较和对比 Elastic Agent 和 Logstash 的描述功能。
10. Logstash 的集中收集能力如何支持 Elastic Agent 的威胁检测能力？

print(chunked_documents[29]['entities'])
'appdynamics', 'apm data', 'azure sentinel', 'microsoft', 'mcafee', 'broadcom', 'cisco', 'dynatrace', 'coveo', 'lucidworks'

{
    "chunk": "...",
    "keyphrases": "...",
    "potential_questions": "...",
    "entities": "..."
}

# 在 embedding_model.py 中定义的嵌入模型
embedder = EmbeddingModel(model_name=HUGGINGFACE_EMBEDDING_MODEL)
cols_to_embed = ['keyphrases', 'potential_questions', 'entities']
embedding_cols = []

for col in cols_to_embed:
    # 处理文本输入
    embedding_col = embedder.embed_documents_text_wise(chunked_documents, text_field=col)
    embedding_cols.append(embedding_col)

# 处理 token 输入
embedding_col = embedder.embed_documents_token_wise(chunked_documents, token_field="chunk")
embedding_cols.append(embedding_col)

embedding_cols = [
    'keyphrases_embedding',
    'potential_questions_embedding',
    'entities_embedding',
    'chunk_embedding'
]

combination_weights = [
    0.1,
    0.15,
    0.05,
    0.7
]

from tqdm import tqdm

def combine_embeddings(objects, embedding_cols, combination_weights, primary_embedding='primary_embedding'):
    # 确保权重数量与嵌入列数量匹配
    assert len(embedding_cols) == len(combination_weights), "嵌入列数量必须与权重数量匹配"
    
    # 归一化权重使其总和为 1
    weights = np.array(combination_weights) / np.sum(combination_weights)
    
    for obj in tqdm(objects, desc="Combining embeddings"):
        # 初始化复合嵌入
        combined = np.zeros_like(obj[embedding_cols[0]])
        
        # 计算加权和
        for col, weight in zip(embedding_cols, weights):
            combined += weight * np.array(obj[col])
        
        # 将新的复合嵌入添加到对象中
        obj.update({primary_embedding: combined.tolist()})
        
        # 删除原始嵌入列
        for col in embedding_cols:
            obj.pop(col, None)

combine_embeddings(chunked_documents, embedding_cols, combination_weights)

{
    'id_': '7fe71686-5cd0-4831-9e79-998c6dbeae0c',
    'chunk': [2312, 14613, ...],
    'original_text': 'if an emerging growth company, indicate by check mark if the registrant has elected not to use the extended ...',
    'chunk_index': 3,
    'chunk_token_count': 399,
    'metadata': {
        'page_label': '3',
        'file_name': 'Elastic_NV_Annual-Report-Fiscal-Year-2023.pdf',
        ...
    },
    'keyphrases': 'sep cl unk\ncheck mark registrant\ncl unk indicate\nunk indicate check\nindicate check mark\nprincipal executive office\naccelerate filer unk\ncompany unk emerge\nunk emerge growth\nemerge growth company',
    'potential_questions': '1. What are the different types of registrant statuses mentioned in the document?\n2. Under what section of the Sarbanes-Oxley Act must registrants file a report on the effectiveness of their internal ...',
    'entities': 'the effectiveness of\nsection 13\nSEP\nUNK\nsection 21e\n1934\n1933\nu. s. c.\nsection 404\nsection 12\nal',
    'primary_embedding': [-0.3946287803351879, -0.17586839850991964, ...]
}

# 初始化 Elasticsearch
ELASTIC_CLOUD_ID = os.environ.get('ELASTIC_CLOUD_ID')
ELASTIC_USERNAME = os.environ.get('ELASTIC_USERNAME')
ELASTIC_PASSWORD = os.environ.get('ELASTIC_PASSWORD')
ELASTIC_CLOUD_AUTH = (ELASTIC_USERNAME, ELASTIC_PASSWORD)

es_bulk_indexer = ESBulkIndexer(cloud_id=ELASTIC_CLOUD_ID, credentials=ELASTIC_CLOUD_AUTH)
es_query_maker = ESQueryMaker(cloud_id=ELASTIC_CLOUD_ID, credentials=ELASTIC_CLOUD_AUTH)

# 定义索引名称
index_name = os.environ.get('ELASTIC_INDEX_NAME')

# 创建索引并批量上传
index_exists = es_bulk_indexer.check_index_existence(index_name=index_name)
if not index_exists:
    logger.info(f"Creating new index: {index_name}")
    es_bulk_indexer.create_es_index(es_configuration=BASIC_CONFIG, index_name=index_name)

success_count = es_bulk_indexer.bulk_upload_documents(
    index_name=index_name,
    documents=chunked_documents,
    id_col='id_',
    batch_size=32
)