RAGViz: Diagnose and Visualize Retrieval-Augmented Generation

TL;DR

RAGViz is an open-source diagnostic tool that visualizes token and document-level attention in retrieval-augmented generation models, enabling better understanding and comparison of generated outputs grounded in retrieved documents.

Contribution

It introduces RAGViz, a novel visualization toolkit for analyzing retrieval-augmented generation models, with efficient operation and customizable components.

Findings

Provides token and document-level attention visualization

Enables comparison of generation with different context documents

Operates efficiently with median query time of 5 seconds

Abstract

Retrieval-augmented generation (RAG) combines knowledge from domain-specific sources into large language models to ground answer generation. Current RAG systems lack customizable visibility on the context documents and the model's attentiveness towards such documents. We propose RAGViz, a RAG diagnosis tool that visualizes the attentiveness of the generated tokens in retrieved documents. With a built-in user interface, retrieval index, and Large Language Model (LLM) backbone, RAGViz provides two main functionalities: (1) token and document-level attention visualization, and (2) generation comparison upon context document addition and removal. As an open-source toolkit, RAGViz can be easily hosted with a custom embedding model and HuggingFace-supported LLM backbone. Using a hybrid ANN (Approximate Nearest Neighbor) index, memory-efficient LLM inference tool, and custom context snippet method, RAGViz operates efficiently with a median query time of about 5 seconds on a moderate GPU node. Our code is available at https://github.com/cxcscmu/RAGViz. A demo video of RAGViz can be found at https://youtu.be/cTAbuTu6ur4.