计算机前沿技术-人工智能算法-大语言模型-最新研究进展-2024-10-27

---

目录

---

1. Large Language Model-based Augmentation for Imbalanced Node Classification on Text-Attributed Graphs

Authors: Leyao Wang, Yu Wang, Bo Ni, Yuying Zhao, Tyler Derr
https://arxiv.org/abs/2410.16882

基于大型语言模型的文本属性图不平衡节点分类数据增强

摘要

本文研究了图上的节点分类问题，特别是面对类别不平衡的挑战。尽管已有多种数据中心的解决方案被提出，但它们并没有专注于文本属性图(TAGs)，因而忽视了利用文本特征中的丰富语义来提升少数类别节点分类的潜力。鉴于此，作者探索了利用大型语言模型(LLMs)在文本空间增强图数据的可能性，提出了一种新颖的方法LA-TAG（基于LLM的文本属性图数据增强），该方法通过提示LLM生成基于图现有节点文本的合成文本。此外，为了将这些合成的文本属性节点整合到图中，引入了一个基于文本的链接预测器来连接合成节点与现有节点。实验结果表明，LA-TAG框架在多个数据集和评估指标上显著优于传统的非文本数据增强策略和特定的节点不平衡解决方案，突出了使用LLM解决TAGs上不平衡问题的可能性。

研究背景

图表示学习在多个领域都非常重要，节点分类是其中的一个基本任务。然而，节点分类经常遇到类别不平衡问题，多数节点往往会主导预测结果，导致对少数节点的结果产生偏差，这在现实世界的应用中可能会带来风险。例如，在假账户检测中，训练模型大多基于正常用户，而少数机器人用户则容易遗漏假账户。类似地，在在线社交网络中，有自杀倾向的个体往往构成少数类别，导致检测和预防覆盖不足。

问题与挑战

现有的方法主要关注传统图，其中节点特征被限制为浅层嵌入，例如，对于文本属性，通常只使用词袋（BOW）特征化。这些方法未能捕捉文本属性中的上下文化语义，导致在基于文本的节点分类任务（如异常检测和自杀识别）中表现不佳。因此，考虑文本特征的文本属性图(TAGs)可以为解决类别不平衡的节点分类问题提供一个可行的解决方案。

如何解决

作者提出了一个名为LA-TAG的框架，它结合了基于LLM的数据增强和基于文本的链接预测器。该框架首先使用LLM生成少数类别的额外文本属性节点，然后使用预训练的语言模型将新合成的文本编码为能够理解文本语义的上下文化深度嵌入。接着，基于文本的链接预测器在新旧节点之间生成连接，保持图的原始几何结构。

创新点

• 首次提出利用LLM进行数据增强来解决TAGs中的不平衡节点分类问题。
• 开发了一个新颖的框架，将基于LLM的数据增强与基于文本的链接预测器相结合，专门针对TAGs。
• 在多个数据集上进行了广泛的评估，包括基线比较、消融研究以及对不同不平衡比例的敏感性分析。

算法模型

LA-TAG框架包括两个主要组件：基于LLM的数据增强和基于文本的链接预测器。数据增强部分，LLM被用来生成额外的文本属性节点，模仿传统的非文本数据增强策略，如上采样、SMOTE和Mixup。链接预测器部分，使用MLP模型来预测节点对之间的链接，基于文本表示来训练，然后将合成的文本数据输入链接预测器，为新节点生成边。

实验效果

实验部分，作者在五个不同的数据集上评估了LA-TAG方法，包括Cora、PubMed、Photo、Computer和Children，涵盖了引用和电子商务领域。评估指标包括整体类别的平均准确率（ACC）、所有类别的平均F1-macro得分（F1）以及多数类别和少数类别平均准确率之间的差异（Diff）。实验结果显示，LA-TAG在整体准确率和宏观F1得分上均优于所有三个基线，同时缩小了多数类别和少数类别之间的平均准确率差异。特别是在Photo、Computer和Children数据集上，与Cora和PubMed相比，改进更为显著。

重要数据与结论

• 在Cora数据集上，LA-TAG的准确率最高可达75.66%，F1得分为74.30%，准确率差异为13.25%。
• 在PubMed数据集上，准确率最高可达76.20%，F1得分为74.35%，准确率差异为12.96%。
• 在Photo数据集上，准确率最高可达66.17%，F1得分为63.62%，准确率差异为11.94%。
• 在Computer数据集上，准确率最高可达64.66%，F1得分为56.36%，准确率差异为23.94%。
• 在Children数据集上，准确率最高可达24.54%，F1得分为22.41%，准确率差异为37.51%。

这些结果表明，LA-TAG在处理文本属性图的类别不平衡节点分类问题上具有显著的优势。

推荐阅读指数: 4.5

2. SleepCoT: A Lightweight Personalized Sleep Health Model via Chain-of-Thought Distillation

Authors: Huimin Zheng, Xiaofeng Xing, Xiangmin Xu
https://arxiv.org/abs/2410.16924

*SleepCoT：一种轻量级个性化睡眠健康模型，通过思维链蒸馏实现

摘要

本文提出了一种新的个性化睡眠健康管理模式，该模式使用少量样本的思维链（Chain-of-Thought，CoT）蒸馏技术，使得小型语言模型（参数少于20亿）在专业健康领域的表现能够与大型语言模型（LLMs）相媲美。这种方法同时从更大的模型中提取问题解决策略、长尾专家知识和个性化推荐能力，并将它们整合到更高效、更紧凑的模型中。与现有系统不同，我们的方法提供了三个关键功能：生成个性化的睡眠健康建议、支持用户特定的后续查询以及回答特定领域的知识问题。我们专注于睡眠健康，因为它可以通过可穿戴设备进行测量，并且对整体健康有着重要影响。我们的实验设置涉及使用GPT-4o进行数据合成、Qwen-max创建指令集以及Qwen2.5 1.5B进行模型蒸馏，与基线小型模型相比，在惩罚、推理和知识应用方面显示出显著改进。使用100个模拟的睡眠报告和1000个特定领域的问题的实验表明，我们的模型在保持实际部署效率的同时实现了与大型模型相当的性能。这项研究不仅推进了AI驱动的健康管理，还提供了一种在资源受限环境中利用LLM能力的新方法，可能增强个性化医疗解决方案的可及性。

研究背景

随着可穿戴设备的快速普及，个人健康数据收集进入了一个新的时代。这些设备持续收集大量的生理数据，作为个人健康状况的外部指标。心率变异性、睡眠模式、体力活动水平等指标提供了一个人的健康状态的全面视图，为个人健康趋势和潜在问题提供了前所未有的洞察。这些设备产生的数据量和多样性为个性化医疗领域带来了机遇和挑战。一方面，这些数据的洪流允许对健康模式进行更细致和个性化的理解。通过分析这些庞大的数据集，可以识别健康状态的微妙变化，可能发现各种状况的早期预警信号或跟踪生活方式变化的有效性。大型语言模型（LLMs）的快速发展，如GPT-4o、Claude 3.5 Sonnet和Qwen-max，显著推进了个性化健康管理和特定领域知识应用的处理。这些最先进的模型能够根据可穿戴设备提供的生理信号，如心率变异性和睡眠模式，生成个性化的生活方式建议。通过利用复杂的推理和专家知识，LLMs可以提供量身定制的建议，在各种情境中增强用户的健康管理。然而，尽管这些模型表现出色，但在实际应用中，特别是在资源受限的环境中，它们面临着显著的限制。首先，这些模型的巨大规模带来了挑战。这些模型拥有数十亿参数，需要大量的计算资源，使得它们在智能手机或可穿戴产品等边缘设备上的部署变得不切实际。这限制了它们在日常场景中的可用性，在这些场景中，设备上的实时交互对于有效的个性化健康管理至关重要。此外，与这些模型相关的高计算复杂性导致了能源消耗增加和硬件成本上升，阻碍了这些高级AI驱动解决方案的广泛采用。其次，与大型模型相关的延迟是一个关键瓶颈。虽然这些模型在计算资源丰富的云环境中表现良好，但由于其庞大的规模，它们的响应时间常常无法满足用户的期望。在涉及个人查询和健康咨询的场景中，这种延迟可能会显著降低用户体验，无法满足交互式应用所需的即时反馈要求。寻求实时洞察和指导的用户无法承受这些过大模型引入的延迟，最终影响了AI驱动健康管理的有效性。

问题与挑战

尽管大型模型在个性化健康管理方面表现出色，但它们在资源受限的环境中面临几个挑战：

1. 模型规模和计算需求：大型模型需要大量的计算资源，不适合在边缘设备上部署。
2. 响应时间：大型模型的响应时间往往无法满足用户对即时反馈的需求。
3. 个性化和领域特定知识的处理：小型模型在处理长尾知识问题和生成高度个性化建议方面表现不佳。

如何解决

为了解决这些挑战，作者提出了SleepCoT，这是一个轻量级的个性化睡眠健康模型，利用思维链（CoT）蒸馏技术。该方法从大型LLMs中提取推理能力、专家知识和推荐策略，并将它们整合到更紧凑的模型中。SleepCoT在较小的模型中捕捉了大型模型的能力，实现了高效的设备内交互，而不影响个性化健康管理的质量。这种方法集中在四个关键领域：数据合成、长尾知识处理、个性化问答和个性化建议生成，每个领域都旨在克服小型模型在健康管理背景下的具体限制。

创新点

• 个性化睡眠健康管理：专注于睡眠健康，这是一个可以通过可穿戴设备测量并影响整体健康的领域。
• 少量样本CoT蒸馏：使用少量样本的思维链蒸馏技术，使得小型模型能够模仿大型模型的性能。
• 三个关键功能：生成个性化睡眠健康建议、支持用户特定的后续查询以及回答特定领域的知识问题。
• 数据合成：利用GPT-4o生成模拟真实世界睡眠健康场景的综合数据集，为蒸馏过程提供多样化和强大的训练基础。
• 长尾知识处理和个性化问答：通过CoT蒸馏，SleepCoT从大型模型中捕获并保留了关键的知识组成部分，显著提高了小型模型对个性化、领域特定问题的回答能力。
• 个性化建议生成：通过蒸馏大型模型的推荐策略，SleepCoT能够在不需要大量计算资源的情况下实时提供有意义的个性化建议。

算法模型

SleepCoT的整体框架包括三个主要组件：数据获取、个性化建议生成和用户问答（包括个性化和领域特定知识查询）。使用GPT-4o进行数据合成，Qwen-max创建指令集，以及Qwen2.5 1.5B进行模型蒸馏。模型使用LoRA（Low-Rank Adaptation）方法进行微调，以确保学生模型（Qwen2.5-1.5B）和（Qwen2.5-0.5B）从教师模型（Qwen-max）中有效适应，同时保持计算效率。

实验效果

实验部分，作者在100个模拟的睡眠报告和1000个特定领域的问题的实验中测试了SleepCoT模型。结果表明，SleepCoT在保持实际部署效率的同时实现了与大型模型相当的性能。具体来说，SleepCoT-1.5B在所有四个评估维度上与Qwen-max表现相当，而SleepCoT-0.5B与SleepCoT-1.5B相比略有性能差距。通过调整数据比例的实验表明，随着个性化问答指令数量的增加，所有四个维度的得分都有明显提高。

重要数据与结论

• Qwen-max：在睡眠领域的专业知识不仅足够，而且在各个评估指标上表现良好，表明该模型对睡眠相关知识有深入的理解，能够以高度的准确性回答问题。
• SleepCoT-1.5B：与Qwen-max在所有四个评估维度上表现相当，证明了小型模型在经过蒸馏后能够保留大型模型的专业知识和推理能力。
• SleepCoT-0.5B：与SleepCoT-1.5B相比，在准确性维度上有所落后，表明即使经过微调，0.5B模型在利用长尾知识方面仍然不如1.5B模型。

推荐阅读指数: 3.5

3. Revealing Hidden Bias in AI: Lessons from Large Language Models

Authors: Django Beatty, Kritsada Masanthia, Teepakorn Kaphol, Niphan Sethi
https://arxiv.org/abs/2410.16927

揭示人工智能中的隐藏偏见：大型语言模型的教训

摘要

随着大型语言模型（LLMs）在招聘流程中的作用日益增强，对AI引起的偏见问题的关注也在增加。本研究检查了由Claude 3.5 Sonnet、GPT-4o、Gemini 1.5和Llama 3.1 405B生成的候选人面试报告中的偏见，重点关注性别、种族和年龄等特征。我们评估了基于LLM的匿名化技术在减少这些偏见方面的有效性。发现表明，尽管匿名化减少了某些偏见——尤其是性别偏见——但其有效性因模型和偏见类型而异。值得注意的是，Llama 3.1 405B表现出的总体偏见最低。此外，我们比较匿名化和非匿名化数据的方法揭示了一种评估LLMs固有偏见的新方法，这超出了招聘应用的范围。本研究强调了仔细选择LLM的重要性，并提出了最小化AI应用中偏见的最佳实践，促进了公平性和包容性。

研究背景

大型语言模型（LLMs）在招聘中的采用正在迅速增加，组织利用AI提高招聘流程的效率。尽管这些高级模型如Claude 3.5 Sonnet、GPT-4o、Gemini 1.5和Llama 3.1 405B被用于生成候选人报告、分析简历和制定面试问题，但人们对AI输出中固有偏见的担忧也在增加，这些偏见可能导致不公平的招聘实践，并基于性别、种族、年龄和其他个人特征延续歧视。解决这些偏见至关重要，以确保AI驱动的招聘工具促进公平和多样性，而不是加剧现有的不平等。招聘中的偏见不仅破坏了道德标准，还带来了战略风险，可能限制劳动力多样性，并使组织面临法律和声誉上的后果。

问题与挑战

在AI驱动的招聘过程中，偏见问题是一个主要挑战，尤其是在性别、种族和年龄等方面。这些偏见可能导致不公平的招聘实践，并延续歧视。此外，偏见的存在还可能限制劳动力的多样性，并使组织面临法律和声誉上的风险。

如何解决

为了解决这些挑战，研究者们采取了以下步骤：

1. 系统化检查：系统化地检查LLM生成的候选人面试报告中存在的偏见。
2. 评估匿名化技术：评估基于LLM的匿名化技术在减轻这些偏见方面的有效性。
3. 比较分析：通过比较匿名化和非匿名化数据，提供了一种新的方法来揭示LLMs中的固有偏见。

创新点

• 多模型比较：研究比较了四种不同的LLMs在生成候选人面试报告时的偏见表现。
• 匿名化技术的有效性评估：研究评估了匿名化技术在减少偏见方面的有效性，特别是在性别偏见方面。
• 新方法提出：提出了一种新的方法来评估LLMs的固有偏见，这种方法超出了招聘应用的范围，可能对其他领域也有影响。

算法模型

研究中使用了四种不同的LLMs：

• Claude 3.5 Sonnet：由Anthropic开发，专注于安全和道德的AI使用，擅长文本摘要和上下文相关内容生成。
• GPT-4o：OpenAI的GPT系列的高级版本，以多语言生成和处理复杂任务而闻名。
• Gemini 1.5：来自Google的DeepMind，专门用于多模态任务，有效理解并生成跨领域内容。
• Llama 3.1 405B：由Meta AI开发，拥有4050亿参数，优化了连贯和上下文丰富的内容生成。

实验效果

实验结果表明：

• 偏见检测：Claude偏见检测器显示，与标准报告相比，匿名化报告的偏见显著减少，尤其是性别偏见。
• 认知偏差检测：标准和匿名化报告的认知偏差相似，概述部分大多没有偏差，而在问题、优势和劣势部分出现了一些“个人化”。
• 非匿名化与匿名化CV的比较：在Gemini和Sonnet中，匿名化CV的偏见显著减少，而在GPT和Llama中，偏见减少或保持不变。

重要数据与结论

• Llama 3.1 405B：在所有模型中总体偏见最低，使其成为生成无偏见报告的强有力候选。
• 匿名化的有效性：匿名化在减少某些偏见，尤其是性别偏见方面有效，但其对其他类型的偏见（如残疾、宗教和政治偏见）的影响较小。

推荐阅读指数: 4.5

4. Order Matters: Exploring Order Sensitivity in Multimodal Large Language Models

Authors: Zhijie Tan, Xu Chu, Weiping Li, Tong Mo
https://arxiv.org/abs/2410.16983

顺序很重要：探索多模态大型语言模型中的顺序敏感性

摘要

多模态大型语言模型（MLLMs）利用包含文本、图像或视频的多模态上下文来解决各种多模态任务。然而，我们发现改变多模态输入的顺序会导致模型性能在高级性能和随机猜测之间波动。这种现象存在于单模态（仅文本或仅图像）和混合模态（图像-文本对）上下文中。此外，我们展示了流行的MLLMs特别关注某些多模态上下文位置，尤其是开头和结尾。利用这种特殊关注，我们将关键视频帧和重要的图像/文本内容放在特殊位置，并提交给MLLM进行推理。这种方法在视频字幕匹配任务中平均提高了14.7%的性能，在视觉问题回答任务中提高了17.8%的性能。此外，我们提出了一个新的指标，位置不变准确率（PIA），以解决MLLM评估中的顺序偏见。我们的研究结果有助于更好地理解多模态上下文学习（MMICL），并提供了在不增加计算成本的情况下增强MLLM性能的实际策略。

研究背景

大型语言模型（LLMs）的成功推动了研究人员探索其在多模态任务中的应用。通过将图像特征与文本嵌入对齐，研究人员将视觉输入整合到LLMs中，发展出多模态大型语言模型（MLLMs）。这些模型继承了LLMs在上下文学习（ICL）中的卓越能力，并在理解和推理图像和视频方面展现出显著的性能。LLMs在上下文学习中的核心问题是提示顺序。已有研究表明，在LLMs中存在更好的提示顺序，当顺序正确时，模型表现良好，而其他顺序则导致性能接近随机。LLMs中提示顺序敏感性的观察引发了两个引人入胜的问题：问题I：“MLLMs中是否存在顺序敏感性？”以及问题II：“哪种顺序对MLLM的性能有益？”

问题与挑战

• 问题I：MLLMs是否也存在顺序敏感性？
• 问题II：哪种顺序对MLLM的性能有益？

如何解决

为了解决这些问题，研究者们设计了以下实验：

1. 实验设计：使用CelebAText-HQ数据集设计了一个4次图像-文本提示，并要求OpenFlamingo生成一个新的面部图像的字幕。
2. 性能影响：展示了改变文本、图像或图像-文本对的提示顺序如何影响模型在图像字幕任务中的性能。
3. 模型间转移性：发现高性能提示不是跨模型可转移的，不同模型对顺序的敏感度不同。
4. 特殊位置关注：发现流行的MLLMs特别关注多模态上下文中的某些位置，尤其是开头和结尾。
5. 性能提升策略：通过在特殊位置放置关键视频帧或重要提示，提高了模型在视频字幕匹配和视觉问题回答任务中的性能。
6. 新评估指标：提出了位置不变准确率（PIA）这一新指标，以消除模型评估中由顺序偏见引起的不公平性。

创新点

• 顺序敏感性的发现：揭示了MLLMs中存在顺序敏感性，并且改变提示顺序可以显著影响MLLMs的性能。
• 特殊位置的性能提升：发现将重要内容放置在多模态上下文的开头或结尾可以提高性能，而无需额外计算。
• 顺序敏感性的应用：证明了在视频理解和视觉问题回答等顺序依赖任务中，利用上下文中的特殊位置可以提高模型性能。
• 新评估指标PIA：提出了PIA这一新指标，以减少模型评估中由顺序偏见引入的偏差。

算法模型

研究中使用了多种MLLMs，包括OpenFlamingo、Qwen-VL-Chat-7B、DeepSeek-VL-7B、IDEFICS9B-Instruct和IDEFICS-v2-8BInstruct。这些模型被用于评估不同顺序对性能的影响，并在视频字幕匹配和视觉问题回答任务中进行了测试。

实验效果

实验结果表明：

• 视频字幕匹配任务：将关键帧放置在视频末尾可以获得最高的模型准确率，与放置在开头和中间相比，平均准确率分别提高了17.78%和16.15%。
• 视觉问题回答任务：将问题相关图像放置在RAG图像集的最后位置可以获得最高的模型准确率，与前三个位置相比，准确率分别提高了21.23%、19.87%和12.36%。
• 位置不变准确率（PIA）：PIA小于准确率和循环评估准确率，因为它的缩放因子Ci Pri，惩罚了MLLM盲目选择Op.A或Op.D的倾向。
  

重要数据与结论

• MLLMs表现出对多模态上下文开头和结尾的偏好，这可能反映了人类的认知行为。
• 在不增加计算成本的情况下，通过在特殊位置放置重要内容可以提高模型性能。
• PIA作为一种新的评估指标，可以减少模型评估中由顺序偏见引起的不公平性。

推荐阅读指数: 4.5

5. SELA: Tree-Search Enhanced LLM Agents for Automated Machine Learning

Authors: Yizhou Chi, Yizhang Lin, Sirui Hong, Duyi Pan, Yaying Fei, Guanghao
Mei, Bangbang Liu, Tianqi Pang, Jacky Kwok, Ceyao Zhang, Bang Liu, Chenglin Wu
https://arxiv.org/abs/2410.17238
代码: https://github.com/geekan/MetaGPT

SELA：树搜索增强的LLM代理，用于自动化机器学习

摘要

本文介绍了一种新颖的自动化机器学习（AutoML）方法，名为SELA（树搜索增强的LLM代理），它利用蒙特卡洛树搜索（MCTS）来优化AutoML过程。SELA框架通过将机器学习管道配置表示为树，使代理能够智能地进行实验，并迭代地完善策略，从而更有效地探索机器学习解决方案空间。这种方法允许SELA基于实验反馈发现最优路径，提高解决方案的整体质量。在20个机器学习数据集上的广泛评估中，我们比较了传统和基于代理的AutoML方法的性能，结果表明SELA在所有数据集上对每个基线实现了65%到80%的胜率。这些结果强调了基于代理的策略在AutoML中的潜力，为解决复杂的机器学习挑战提供了新视角。

研究背景

自动化机器学习（AutoML）是一个快速发展的领域，旨在最小化人为干预，自动化设计可靠的机器学习解决方案。传统的AutoML框架依赖于预定义的搜索空间和例程，主要集中在优化超参数和模型集成上。然而，这种方法在处理多样化和动态数据场景时缺乏适应性，导致在更复杂的环境中表现不佳。此外，传统方法主要关注模型训练，而忽视了机器学习管道中其他关键阶段，如数据预处理和特征工程，限制了这些系统的总体有效性。

问题与挑战

尽管基于LLM的代理在自动化机器学习任务中显示出了希望，但它们在生成多样化和高度优化的工作流程方面常常不足，因为它们的搜索过程仅限于单次尝试。没有迭代细化或探索替代策略的能力，这些解决方案经常收敛于次优结果，即使允许多次尝试。

如何解决

SELA框架通过结合LLM代理的优势和结构化搜索与细化过程，模仿专家解决机器学习问题的方式，解决了上述挑战。SELA将机器学习问题的搜索空间表示为树，其中每个分支代表一个潜在的解决方案路径。通过应用MCTS作为核心决策引擎，利用其平衡探索（测试新策略）和利用（改进已知良好策略）的能力，允许代理有效地探索大型决策空间，收集和处理实验结果，并智能选择下一个有前景的配置进行测试。通过这种实验和细化的迭代循环，SELA逐步改进其解决方案。

创新点

1. 反馈驱动方法：为LLM代理提供了一种迭代探索机器学习配置的方法，通过多轮实验优化解决方案。
2. 树结构搜索空间：使用MCTS导航树结构搜索空间，通过反馈自适应地识别高性能管道。
3. 与传统AutoML的比较：突出了基于代理方法在机器学习中的灵活性和增强性能的潜力。

算法模型

SELA框架由三个关键组件构成：基于LLM的洞察提议者、使用MCTS的搜索模块，以及作为实验执行者的LLM代理。框架首先由LLM生成针对机器学习工作流程不同阶段的多样化方法，然后将搜索空间组织成树结构，并使用MCTS探索有前景的路径。在每个循环中，选定的路径被传递给LLM代理，后者将配置转化为可执行的管道，计划、编码并执行实验，并将结果反馈以细化未来的搜索。

实验效果

在20个不同的机器学习数据集上的实验结果表明，SELA在所有数据集上对每个基线实现了65%到80%的胜率，证明了SELA在各种机器学习任务中的优越性能。SELA在产生最高预测结果的数量上表现出色，并且在不同数据集上的平均归一化得分（NS）和平均最佳排名上均优于其他框架。

重要数据与结论

• SELA在所有框架中平均归一化得分（NS）和平均最佳排名最高。
• SELA在“Top 1”列中表现出色，表明其在所有数据集中产生最高预测结果的数量最多。
• SELA在65-80%的数据集上击败了每个竞争方法，显示了其在多样化数据集上的强大性能和适应性。

推荐阅读指数: 5

"""BM25 retriever."""
from typing import Callable, Optional

from llama_index.core import VectorStoreIndex
from llama_index.core.callbacks.base import CallbackManager
from llama_index.core.constants import DEFAULT_SIMILARITY_TOP_K
from llama_index.core.schema import BaseNode, IndexNode
from llama_index.retrievers.bm25 import BM25Retriever
from rank_bm25 import BM25Okapi


class DynamicBM25Retriever(BM25Retriever):
    """BM25 retriever."""

    def __init__(
        self,
        nodes: list[BaseNode],
        tokenizer: Optional[Callable[[str], list[str]]] = None,
        similarity_top_k: int = DEFAULT_SIMILARITY_TOP_K,
        callback_manager: Optional[CallbackManager] = None,
        objects: Optional[list[IndexNode]] = None,
        object_map: Optional[dict] = None,
        verbose: bool = False,
        index: VectorStoreIndex = None,
    ) -> None:
        super().__init__(
            nodes=nodes,
            tokenizer=tokenizer,
            similarity_top_k=similarity_top_k,
            callback_manager=callback_manager,
            object_map=object_map,
            objects=objects,
            verbose=verbose,
        )
        self._index = index

    def add_nodes(self, nodes: list[BaseNode], **kwargs) -> None:
        """Support add nodes."""
        self._nodes.extend(nodes)
        self._corpus = [self._tokenizer(node.get_content()) for node in self._nodes]
        self.bm25 = BM25Okapi(self._corpus)

        if self._index:
            self._index.insert_nodes(nodes, **kwargs)

    def persist(self, persist_dir: str, **kwargs) -> None:
        """Support persist."""
        if self._index:
            self._index.storage_context.persist(persist_dir)

"""FAISS retriever."""

from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.schema import BaseNode


class FAISSRetriever(VectorIndexRetriever):
    """FAISS retriever."""

    def add_nodes(self, nodes: list[BaseNode], **kwargs) -> None:
        """Support add nodes."""
        self._index.insert_nodes(nodes, **kwargs)

    def persist(self, persist_dir: str, **kwargs) -> None:
        """Support persist."""
        self._index.storage_context.persist(persist_dir)

---

后记

如果觉得我的博客对您有用，欢迎 打赏 支持！三连击 (点赞、收藏、关注和评论) 不迷路，我将持续为您带来计算机人工智能前沿技术(尤其是AI相关的大语言模型，深度学习和计算机视觉相关方向)最新学术论文及工程实践方面的内容分享，助力您更快更准更系统地了解 AI前沿技术。