跳至主要内容

结合 RAG 的 Llama 4:含演示项目的指南

学习如何使用 Llama 4 构建检索增强生成(RAG)流程,创建一个简单的 Web 应用程序。
· 12分钟 阅读

使用 AI 探索

Llama 4 Scout 宣称具备 1000 万 tokens 的超大上下文窗口,但其训练时的最大输入规模仅为 256k tokens。这意味着在更大的输入下性能可能会下降。为避免这种情况,我们可以将 Llama 4 与检索增强生成(RAG)流程结合使用。

在本教程中,我将逐步讲解如何使用 LangChain 生态构建 RAG 流程,并创建一个允许用户上传文档并就其内容提问的 Web 应用。

我们通过每周五发送的免费通讯 The Median 为读者带来最新的 AI 资讯,精炼梳理本周要闻。订阅即可每周用几分钟保持敏锐:

为何将 Llama 4 与 RAG 搭配?

Llama 4 提供 1000 万 tokens 的上下文窗口,意味着我们可以喂入 100 本书并获得具有上下文的回答。然而,即便有这一令人印象深刻的特性,仍有充分理由将 Llama 4 与 RAG 搭配。原因如下:

1. 大输入下的性能限制

虽然 Llama 4 支持超大上下文窗口,但其训练时的最大输入仅为 256k tokens。当输入接近或超过该阈值时,模型性能可能下降。RAG 通过仅检索最相关的信息来缓解这一问题,确保输入精炼而聚焦。

2. 提高效率与聚焦性

RAG 能从 向量数据库检索高度相关的信息,使 Llama 4 处理更小且更有针对性的输入,从而同时提升效率与准确性。

3. 节省成本

用 Llama 4 的完整上下文窗口处理海量输入会带来高昂的计算成本。通过 RAG,仅检索和处理最相关的数据,可显著降低计算负载与相关费用。

构建 Llama 4 的 RAG 流程

本教程将逐步讲解如何构建一个简单的检索增强生成流程,使您可以加载 .docx 文件并就其内容提问。

1. 环境准备

首先,我们需要创建一个 GroqCloud API 密钥,并将其保存为环境变量。接着,我们将安装用于加载、切分、向量化、检索上下文以及生成回答所需的全部 Python 包。

%%capture
%pip install langchain
%pip install langchain-community 
%pip install langchainhub 
%pip install langchain-chroma 
%pip install langchain-groq
%pip install langchain-huggingface
%pip install unstructured[docx]

2. 初始化 LLM 与嵌入模型

我们将使用 Groq API 设置模型对象,提供模型名称与 API 密钥。同时,从 Hugging Face 下载嵌入模型并将其作为我们的嵌入模型加载。

from langchain_groq import ChatGroq
from langchain_huggingface import HuggingFaceEmbeddings

llm = ChatGroq(model="meta-llama/llama-4-scout-17b-16e-instruct", api_key=groq_api_key)

embed_model = HuggingFaceEmbeddings(model_name="mixedbread-ai/mxbai-embed-large-v1")

3. 加载并切分数据

我们将从根目录加载 .docx 文件,并将其切分为每段 1000 个字符的块。

from langchain_community.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

# Initialize the text splitter
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=100,
    separators=["\n\n", "\n"]
)

# Load the .docx files
loader = DirectoryLoader("./", glob="*.docx", use_multithreading=True)
documents = loader.load()

# Split the documents into chunks
chunks = text_splitter.split_documents(documents)

# Print the number of chunks
print(len(chunks))
29

4. 构建并填充向量库

初始化 Chroma 向量库,并将文本块提供给向量数据库。文本会在存储前被转换为嵌入。之后,我们将运行一次相似度搜索,以测试向量库是否已正确填充。

from langchain_chroma import Chroma

vectorstore = Chroma.from_documents(
    documents=chunks,
    embedding=embed_model,
    persist_directory="./Vectordb",
)

query = "What is this tutorial about?"
docs = vectorstore.similarity_search(query)
print(docs[0].page_content)
Learn how to Fine-tune Stable Diffusion XL with DreamBooth and LoRA on your personal images. 

Let's try another prompt:

Prompt:

5. 创建 RAG 流程

接下来,我们将把向量库转换为检索器,并为 RAG 流程创建提示模板。

# Create retriever
retriever = vectorstore.as_retriever()

# Import PromptTemplate
from langchain_core.prompts import PromptTemplate

# Define a clearer, more professional prompt template
template = """You are an expert assistant tasked with answering questions based on the provided documents.
Use only the given context to generate your answer.
If the answer cannot be found in the context, clearly state that you do not know.
Be detailed and precise in your response, but avoid mentioning or referencing the context itself.

Context:
{context}

Question:
{question}

Answer:"""

# Create the PromptTemplate
rag_prompt = PromptTemplate.from_template(template)

最后,我们将创建 RAG 链,在 RAG 提示中同时提供上下文与问题,传入 Llama 4 模型,然后生成干净的回答。

from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough

rag_chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | rag_prompt
    | llm
    | StrOutputParser()
)

现在我们来就该文档向 RAG 链提问。

from IPython.display import display, Markdown

response = rag_chain.invoke("What this tutorial about?")
Markdown(response)
This tutorial is about setting up and using the Janus project, specifically Janus Pro, a multimodal model that can understand images and generate images from text prompts, and building a local solution to use the model privately on a laptop GPU. It covers learning about the Janus Series, setting up the Janus project, building a Docker container to run the model locally, and testing its capabilities with various image and text prompts.

如您所见,回答具有很强的准确性与上下文意识。如果运行上述代码遇到问题,请参考此 DataLab 工作簿: Llama 4 RAG pipeline。

构建 Llama 4 Docx 聊天应用

本节中,我们将把上文的 RAG 流程代码改造成一个使用 Gradio 包的可用聊天应用。首先,您需要通过以下 pip 命令安装 Gradio:

$ pip install gradio

该应用设计简洁友好,支持上传单个 .docx 文件、多个 .docx 文件,或包含 .docx 文件的 .zip 压缩包。随后,用户可以就所有已上传文件的内容进行提问,以快速高效地检索信息。

应用特性包括:

  • 文档加载: 应用可直接从上传文件中提取并加载文档。
  • 文本切分: 将文档切分为便于处理与检索的块。
  • 向量库集成: 与其使用 ChromaDB,本应用采用内存向量库管理文档嵌入。
  • RAG 流程: 结合文档检索与语言模型,基于上传内容提供准确回答。
  • 用户界面: 提供聊天界面,用户可就提供的文档提问并获得详细回复。
  • 错误处理: 对不支持的文件类型和上传问题提供反馈,提升用户体验。
  • 重置功能: 用户可重置应用状态,重新上传新文档开始对话。

我们把所有内容整合如下:

main.py:

# ========== Standard Library ==========
import os
import tempfile
import zipfile
from typing import List, Optional, Tuple, Union
import collections


# ========== Third-Party Libraries ==========
import gradio as gr
from groq import Groq
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import DirectoryLoader, UnstructuredFileLoader
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_groq import ChatGroq
from langchain_huggingface import HuggingFaceEmbeddings

# ========== Configs ==========
TITLE = """<h1 align="center">🗨️🦙 Llama 4 Docx Chatter</h1>"""
AVATAR_IMAGES = (
    None,
    "./logo.png",
)

# Acceptable file extensions
TEXT_EXTENSIONS = [".docx", ".zip"]

# ========== Models & Clients ==========
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
client = Groq(api_key=GROQ_API_KEY)
llm = ChatGroq(model="meta-llama/llama-4-scout-17b-16e-instruct", api_key=GROQ_API_KEY)
embed_model = HuggingFaceEmbeddings(model_name="mixedbread-ai/mxbai-embed-large-v1")

# ========== Core Components ==========
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=100,
    separators=["\n\n", "\n"],
)

rag_template = """You are an expert assistant tasked with answering questions based on the provided documents.
Use only the given context to generate your answer.
If the answer cannot be found in the context, clearly state that you do not know.
Be detailed and precise in your response, but avoid mentioning or referencing the context itself.
Context:
{context}
Question:
{question}
Answer:"""
rag_prompt = PromptTemplate.from_template(rag_template)


# ========== App State ==========
class AppState:
    vectorstore: Optional[InMemoryVectorStore] = None
    rag_chain = None


state = AppState()

# ========== Utility Functions ==========


def load_documents_from_files(files: List[str]) -> List:
    """Load documents from uploaded files directly without moving."""
    all_documents = []

    # Temporary directory if ZIP needs extraction
    with tempfile.TemporaryDirectory() as temp_dir:
        for file_path in files:
            ext = os.path.splitext(file_path)[1].lower()

            if ext == ".zip":
                # Extract ZIP inside temp_dir
                with zipfile.ZipFile(file_path, "r") as zip_ref:
                    zip_ref.extractall(temp_dir)

                # Load all docx from extracted zip
                loader = DirectoryLoader(
                    path=temp_dir,
                    glob="**/*.docx",
                    use_multithreading=True,
                )
                docs = loader.load()
                all_documents.extend(docs)

            elif ext == ".docx":
                # Load single docx directly
                loader = UnstructuredFileLoader(file_path)
                docs = loader.load()
                all_documents.extend(docs)

    return all_documents


def get_last_user_message(chatbot: List[Union[gr.ChatMessage, dict]]) -> Optional[str]:
    """Get last user prompt."""
    for message in reversed(chatbot):
        content = (
            message.get("content") if isinstance(message, dict) else message.content
        )
        if (
            message.get("role") if isinstance(message, dict) else message.role
        ) == "user":
            return content
    return None


# ========== Main Logic ==========




def upload_files(
    files: Optional[List[str]], chatbot: List[Union[gr.ChatMessage, dict]]
):
    """Handle file upload - .docx or .zip containing docx."""
    if not files:
        return chatbot

    file_summaries = []  # <-- Collect formatted file/folder info
    documents = []

    with tempfile.TemporaryDirectory() as temp_dir:
        for file_path in files:
            filename = os.path.basename(file_path)
            ext = os.path.splitext(file_path)[1].lower()

            if ext == ".zip":
                file_summaries.append(f"📦 **{filename}** (ZIP file) contains:")
                try:
                    with zipfile.ZipFile(file_path, "r") as zip_ref:
                        zip_ref.extractall(temp_dir)
                        zip_contents = zip_ref.namelist()

                        # Group files by folder
                        folder_map = collections.defaultdict(list)
                        for item in zip_contents:
                            if item.endswith("/"):
                                continue  # skip folder entries themselves
                            folder = os.path.dirname(item)
                            file_name = os.path.basename(item)
                            folder_map[folder].append(file_name)

                        # Format nicely
                        for folder, files_in_folder in folder_map.items():
                            if folder:
                                file_summaries.append(f"📂 {folder}/")
                            else:
                                file_summaries.append(f"📄 (root)")
                            for f in files_in_folder:
                                file_summaries.append(f"   - {f}")

                    # Load docx files extracted from ZIP
                    loader = DirectoryLoader(
                        path=temp_dir,
                        glob="**/*.docx",
                        use_multithreading=True,
                    )
                    docs = loader.load()
                    documents.extend(docs)

                except zipfile.BadZipFile:
                    chatbot.append(
                        gr.ChatMessage(
                            role="assistant",
                            content=f"❌ Failed to open ZIP file: {filename}",
                        )
                    )

            elif ext == ".docx":
                file_summaries.append(f"📄 **{filename}**")
                loader = UnstructuredFileLoader(file_path)
                docs = loader.load()
                documents.extend(docs)

            else:
                file_summaries.append(f"❌ Unsupported file type: {filename}")

    if not documents:
        chatbot.append(
            gr.ChatMessage(
                role="assistant", content="No valid .docx files found in upload."
            )
        )
        return chatbot

    # Split documents
    chunks = text_splitter.split_documents(documents)
    if not chunks:
        chatbot.append(
            gr.ChatMessage(
                role="assistant", content="Failed to split documents into chunks."
            )
        )
        return chatbot

    # Create Vectorstore
    state.vectorstore = InMemoryVectorStore.from_documents(
        documents=chunks,
        embedding=embed_model,
    )
    retriever = state.vectorstore.as_retriever()

    # Build RAG Chain
    state.rag_chain = (
        {"context": retriever, "question": RunnablePassthrough()}
        | rag_prompt
        | llm
        | StrOutputParser()
    )

    # Final display
    chatbot.append(
        gr.ChatMessage(
            role="assistant",
            content="**Uploaded Files:**\n"
            + "\n".join(file_summaries)
            + "\n\n✅ Ready to chat!",
        )
    )
    return chatbot


def user_message(
    text_prompt: str, chatbot: List[Union[gr.ChatMessage, dict]]
) -> Tuple[str, List[Union[gr.ChatMessage, dict]]]:
    """Add user's text input to conversation."""
    if text_prompt.strip():
        chatbot.append(gr.ChatMessage(role="user", content=text_prompt))
    return "", chatbot


def process_query(
    chatbot: List[Union[gr.ChatMessage, dict]],
) -> List[Union[gr.ChatMessage, dict]]:
    """Process user's query through RAG pipeline."""
    prompt = get_last_user_message(chatbot)
    if not prompt:
        chatbot.append(
            gr.ChatMessage(role="assistant", content="Please type a question first.")
        )
        return chatbot

    if state.rag_chain is None:
        chatbot.append(
            gr.ChatMessage(role="assistant", content="Please upload documents first.")
        )
        return chatbot

    chatbot.append(gr.ChatMessage(role="assistant", content="Thinking..."))

    try:
        response = state.rag_chain.invoke(prompt)
        chatbot[-1].content = response
    except Exception as e:
        chatbot[-1].content = f"Error: {str(e)}"

    return chatbot


def reset_app(
    chatbot: List[Union[gr.ChatMessage, dict]],
) -> List[Union[gr.ChatMessage, dict]]:
    """Reset application state."""
    state.vectorstore = None
    state.rag_chain = None
    return [
        gr.ChatMessage(
            role="assistant", content="App reset! Upload new documents to start."
        )
    ]


# ========== UI Layout ==========

with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.HTML(TITLE)
    chatbot = gr.Chatbot(
        label="Llama 4 RAG",
        type="messages",
        bubble_full_width=False,
        avatar_images=AVATAR_IMAGES,
        scale=2,
        height=350,
    )

    with gr.Row(equal_height=True):
        text_prompt = gr.Textbox(
            placeholder="Ask a question...", show_label=False, autofocus=True, scale=28
        )
        send_button = gr.Button(
            value="Send",
            variant="primary",
            scale=1,
            min_width=80,
        )
        upload_button = gr.UploadButton(
            label="Upload",
            file_count="multiple",
            file_types=TEXT_EXTENSIONS,
            scale=1,
            min_width=80,
        )
        reset_button = gr.Button(
            value="Reset",
            variant="stop",
            scale=1,
            min_width=80,
        )

    send_button.click(
        fn=user_message,
        inputs=[text_prompt, chatbot],
        outputs=[text_prompt, chatbot],
        queue=False,
    ).then(fn=process_query, inputs=[chatbot], outputs=[chatbot])

    text_prompt.submit(
        fn=user_message,
        inputs=[text_prompt, chatbot],
        outputs=[text_prompt, chatbot],
        queue=False,
    ).then(fn=process_query, inputs=[chatbot], outputs=[chatbot])

    upload_button.upload(
        fn=upload_files, inputs=[upload_button, chatbot], outputs=[chatbot], queue=False
    )
    reset_button.click(fn=reset_app, inputs=[chatbot], outputs=[chatbot], queue=False)

demo.queue().launch()

测试 Llama 4 Docx 聊天应用

现在我们来运行该 Web 应用并测试其功能。要启动应用,请在终端中运行以下命令:

$ python main.py 

当 Web 服务器启动后,访问以下 URL,或将其复制粘贴到浏览器地址栏以进入 Web 界面: http://127.0.0.1:7860

聊天界面如下所示:

llama 4 rag gradio 接口

我们将上传一个包含 .docx 文件的 .zip 压缩包,在上传过程中,您将看到压缩包内的所有文件夹与文件。

llama 4 rag gradio 接口测试

接下来,我们可以就文档提出具体问题:

llama 4 rag gradio 接口测试

您也可以点击“Reset”按钮,在几秒钟内上传新文件并与之对话。

llama 4 rag gradio 接口测试

在本地测试完成后,您可以轻松将该 Web 应用免费部署到 Hugging Face Space。这里是已部署应用的链接: kingabzpro/Llama-4-RAG 

llama 4 rag gradio 接口测试

来源: kingabzpro/Llama-4-RAG 

代码源、笔记本与配置文件可在 kingabzpro/Llama-4-RAG 存储库中找到。

结语

即便长上下文窗口模型不断涌现,RAG 也不会过时。RAG 仍是从外部资源检索信息、为模型提供额外上下文的高效且具成本效益的方法。

此外,与直接向模型提供大量上下文(需要大量计算资源)相比,RAG 系统在运行成本上更具优势。这使得 RAG 成为众多用例中务实且可扩展的方案。

如需进一步了解 Llama 4,请查看以下资源:

Abid Ali Awan

作者

Abid Ali Awan

作为一名持证的数据科学家,我热衷于利用前沿技术打造创新的机器学习应用。凭借在语音识别、数据分析与报告、MLOps、对话式人工智能以及自然语言处理方面的扎实背景,我不断打磨构建智能系统的能力,力求带来切实影响。除技术专长外,我也擅长沟通,能够将复杂概念提炼为清晰、简明的表述。因此,我成为数据科学领域备受关注的博主,与不断壮大的数据专业人士社区分享洞见与实践经验。目前,我专注于内容创作与编辑,借助大语言模型打造有力且吸引人的内容,帮助企业与个人更好地发挥数据价值。

主题