向HuggingFace社区介绍FedRAG

今天，我很高兴能与HuggingFace社区分享FedRAG。这是我自1月以来一直致力于的项目，并于上个月“发布”（请参阅Medium上的介绍性博客文章）。

什么是FedRAG？

FedRAG是一个用于跨中心化和联邦架构微调RAG系统的框架。它原生支持PyTorch，并支持一些关键集成，最著名的是与HuggingFace的集成，这使得RAG系统可以灵活组装以进行微调。

我们的研讨会论文已被ICML 2025机器学习开源开发（CODEML '25）研讨会录用。 您可以在此处下载我们的终稿论文。

您今天可以通过以下方式安装FedRAG：pip install fed-rag[huggingface]。

FedRAG模式

在本节中，我们提供代码片段来说明FedRAG中的关键使用模式。

组装RAG系统

# Flat imports are supported
from fed_rag import (
  RAGSystem,
  RAGConfig,
  HFSentenceTransformerRetriever,
  UnslothFastModelGenerator,
  QdrantKnowledgeStore
)

knowledge_store = QdrantKnowledgeStore()
generator = UnslothFastModelGenerator(
  "unsloth/gemma-3-4b-it",
)
retriever = HFSentenceTransformerRetriever(
  "nthakur/dragon-plus-query-encoder",
  "nthakur/dragon-plus-context-encoder",
)

# Assemble rag_system
rag_system = RAGSystem(
  knowledge_store=knowledge_store,
  generator=generator,
  retriever=retriever,
  rag_config=RAGConfig(top_k=2)
)

# Executes the typical RAG pipeline
response = rag_system.query("What are tulips?")

生成器和检索器微调

#...keep the code from the previous pattern listing
from fed_rag.trainers import (
  HuggingFaceTrainerForRALT,
  HuggingFaceTrainerForLSR
)
from fed_rag.trainer_managers import (
  HuggingFaceRAGTrainerManager
)
from datasets import Dataset

# Train datasets are examples of (query, 
# response) pairs
train_dataset = Dataset.from_dict(
  {
    "query": [...],
    "response": [...]
  }
)
generator_trainer = HuggingFaceTrainerForRALT(
    rag_system=rag_system,
    train_dataset=train_dataset,
)
retriever_trainer = HuggingFaceTrainerForLSR(
    rag_system=rag_system,
    train_dataset=train_dataset,
)
manager = HuggingFaceRAGTrainerManager(
    mode="retriever",  # can be generator
    retriever_trainer=retriever_trainer,
    generator_trainer=generator_trainer,
)

# Train
train_result = manager.train()

RALT和LSR训练方法在RA-DIT论文 (Lin et al. (2023) 中介绍。

联邦微调

... # Keep code from listings 1 and 2
import flwr as fl  # The FL backend
# fl_task
fl_task = manager.get_federated_task()
# Build fl server and client
server = fl_task.server(
  model=retriever_trainer.model
)
client = fl_task.client(
  model=retriever_trainer.model,
  train_dataset=train_dataset,
)
# Spin up client and server using flwr
fl.start_server(server)
fl.start_client(client)

FedRAG适用于谁？

我们构建FedRAG时考虑到了从业者和研究人员。我们致力于成为事实上的RAG研究框架，使研究人员能够生产可靠且可重现的最新技术。我们与诸如~transformers和~sentence_transformers等工具的深度集成意味着先进的研究可以更快地、更广泛地惠及从业者。

链接

• ⭐️ 在Github上给我们点星
• 📓 官方文档
• 💬 加入我们的Discord