当前位置:AIGC资讯 > AIGC > 正文

16 使用Llama Index提取术语和定义的指南

使用Llama Index提取术语和定义的指南

Llama Index有许多用例(如语义搜索、总结等),这些用例都有很好的文档记录。然而,这并不意味着我们不能将Llama Index应用于非常特定的用例!

在本教程中,我们将介绍使用Llama Index从文本中提取术语和定义的设计过程,同时允许用户稍后查询这些术语。使用Streamlit,我们可以提供一种简单的方法来构建前端,运行和测试所有这些内容,并快速迭代我们的设计。

本教程假设你已经安装了Python3.9+和以下包:

llama-index streamlit

在基本层面上,我们的目标是获取文档中的文本,提取术语和定义,然后为用户提供一种查询这些术语和定义知识库的方法。本教程将介绍Llama Index和Streamlit的功能,并希望为出现的常见问题提供一些有趣的解决方案。

本教程的最终版本可以在这里找到,一个实时的托管演示可以在Huggingface Spaces上找到。

上传文本

第一步是为用户提供一种手动输入文本的方式。让我们使用Streamlit编写一些代码来提供这个界面!使用以下代码并使用streamlit run app.py启动应用程序。

import streamlit as st

st.title("? Llama Index Term Extractor ?")

document_text = st.text_area("Enter raw text")
if st.button("Extract Terms and Definitions") and document_text:
    with st.spinner("Extracting..."):
        extracted_terms = document_text  # this is a placeholder!
    st.write(extracted_terms)

非常简单对吧!但你会注意到应用程序还没有做任何有用的事情。要使用llama_index,我们还需要设置我们的OpenAI LLM。LLM有许多可能的设置,所以我们可以让用户自己找出最佳设置。我们还应该让用户设置将提取术语的提示(这也会帮助我们调试什么效果最好)。

LLM设置

这一步在我们的应用程序中引入了一些标签,将其分成不同的窗格,提供不同的功能。让我们为LLM设置和上传文本创建一个标签:

import os
import streamlit as st

DEFAULT_TERM_STR = (
    "Make a list of terms and definitions that are defined in the context, "
    "with one pair on each line. "
    "If a term is missing it's definition, use your best judgment. "
    "Write each line as as follows:\nTerm: <term> Definition: <definition>"
)

st.title("? Llama Index Term Extractor ?")

setup_tab, upload_tab = st.tabs(["Setup", "Upload/Extract Terms"])

with setup_tab:
    st.subheader("LLM Setup")
    api_key = st.text_input("Enter your OpenAI API key here", type="password")
    llm_name = st.selectbox("Which LLM?", ["gpt-3.5-turbo", "gpt-4"])
    model_temperature = st.slider(
        "LLM Temperature", min_value=0.0, max_value=1.0, step=0.1
    )
    term_extract_str = st.text_area(
        "The query to extract terms and definitions with.",
        value=DEFAULT_TERM_STR,
    )

with upload_tab:
    st.subheader("Extract and Query Definitions")
    document_text = st.text_area("Enter raw text")
    if st.button("Extract Terms and Definitions") and document_text:
        with st.spinner("Extracting..."):
            extracted_terms = document_text  # this is a placeholder!
        st.write(extracted_terms)

现在我们的应用程序有两个标签,这真的有助于组织。你还注意到我添加了一个默认提示来提取术语——你可以在尝试提取一些术语后更改这个提示,这只是我在实验后得出的提示。

说到提取术语,是时候添加一些函数来完成这个任务了!

提取和存储术语

现在我们能够定义LLM设置并输入文本,我们可以尝试使用Llama Index从文本中提取术语!

我们可以添加以下函数来初始化我们的LLM,并使用它从输入文本中提取术语。

from llama_index.core import Document, SummaryIndex, load_index_from_storage
from llama_index.llms.openai import OpenAI
from llama_index.core import Settings

def get_llm(llm_name, model_temperature, api_key, max_tokens=256):
    os.environ["OPENAI_API_KEY"] = api_key
    return OpenAI(
        temperature=model_temperature, model=llm_name, max_tokens=max_tokens
    )

def extract_terms(
    documents, term_extract_str, llm_name, model_temperature, api_key
):
    llm = get_llm(llm_name, model_temperature, api_key, max_tokens=1024)

    temp_index = SummaryIndex.from_documents(
        documents,
    )
    query_engine = temp_index.as_query_engine(
        response_mode="tree_summarize", llm=llm
    )
    terms_definitions = str(query_engine.query(term_extract_str))
    terms_definitions = [
        x
        for x in terms_definitions.split("\n")
        if x and "Term:" in x and "Definition:" in x
    ]
    # parse the text into a dict
    terms_to_definition = {
        x.split("Definition:")[0]
        .split("Term:")[-1]
        .strip(): x.split("Definition:")[-1]
        .strip()
        for x in terms_definitions
    }
    return terms_to_definition

现在,使用新函数,我们终于可以提取我们的术语了!

...
with upload_tab:
    st.subheader("Extract and Query Definitions")
    document_text = st.text_area("Enter raw text")
    if st.button("Extract Terms and Definitions") and document_text:
        with st.spinner("Extracting..."):
            extracted_terms = extract_terms(
                [Document(text=document_text)],
                term_extract_str,
                llm_name,
                model_temperature,
                api_key,
            )
        st.write(extracted_terms)

现在有很多事情在进行,让我们花点时间回顾一下发生了什么。

get_llm() 根据设置标签中的用户配置实例化LLM。根据模型名称,我们需要使用适当的类(OpenAI vs. ChatOpenAI)。

extract_terms() 是所有好东西发生的地方。首先,我们用max_tokens=1024调用get_llm(),因为我们不希望在提取术语和定义时限制模型太多(默认值是256,如果未设置)。然后,我们定义我们的Settings对象,将num_output与我们的max_tokens值对齐,并将块大小设置为不大于输出。当文档被Llama Index索引时,如果它们很大,它们会被分成块(也称为节点),chunk_size设置这些块的大小。

接下来,我们创建一个临时的总结索引并传入我们的llm。总结索引将读取索引中的每一条文本,这对于提取术语非常完美。最后,我们使用我们预定义的查询文本提取术语,使用response_mode="tree_summarize。这种响应模式将从底部向上生成一个总结树,每个父节点总结其子节点。最后,返回树的顶部,其中将包含我们提取的所有术语和定义。

最后,我们进行一些次要的后处理。我们假设模型遵循指令并在每行上放置一个术语/定义对。如果一行缺少Term:或Definition:标签,我们跳过它。然后,我们将其转换为字典以便于存储!

保存提取的术语

现在我们可以提取术语,我们需要将它们放在某个地方,以便稍后查询它们。VectorStoreIndex现在应该是一个完美的选择!此外,我们的应用程序还应该跟踪插入索引中的术语,以便稍后检查它们。使用st.session_state,我们可以将当前术语列表存储在每个用户唯一的会话字典中!

首先,让我们添加一个功能来初始化全局向量索引,并添加另一个函数来插入提取的术语。

from llama_index.core import Settings, VectorStoreIndex

...
if "all_terms" not in st.session_state:
    st.session_state["all_terms"] = DEFAULT_TERMS
...

def insert_terms(terms_to_definition):
    for term, definition in terms_to_definition.items():
        doc = Document(text=f"Term: {term}\nDefinition: {definition}")
        st.session_state["llama_index"].insert(doc)

@st.cache_resource
def initialize_index(llm_name, model_temperature, api_key):
    """Create the VectorStoreIndex object."""
    Settings.llm = get_llm(llm_name, model_temperature, api_key)

    index = VectorStoreIndex([])

    return index, llm

...

with upload_tab:
    st.subheader("Extract and Query Definitions")
    if st.button("Initialize Index and Reset Terms"):
        st.session_state["llama_index"] = initialize_index(
            llm_name, model_temperature, api_key
        )
        st.session_state["all_terms"] = {}

    if "llama_index" in st.session_state:
        st.markdown(
            "Either upload an image/screenshot of a document, or enter the text manually."
        )
        document_text = st.text_area("Or enter raw text")
        if st.button("Extract Terms and Definitions") and (
            uploaded_file or document_text
        ):
            st.session_state["terms"] = {}
            terms_docs = {}
            with st.spinner("Extracting..."):
                terms_docs.update(
                    extract_terms(
                        [Document(text=document_text)],
                        term_extract_str,
                        llm_name,
                        model_temperature,
                        api_key,
                    )
                )
            st.session_state["terms"].update(terms_docs)

        if "terms" in st.session_state and st.session_state["terms"]:
            st.markdown("Extracted terms")
            st.json(st.session_state["terms"])

            if st.button("Insert terms?"):
                with st.spinner("Inserting terms"):
                    insert_terms(st.session_state["terms"])
                st.session_state["all_terms"].update(st.session_state["terms"])
                st.session_state["terms"] = {}
                st.experimental_rerun()

现在你真的开始利用streamlit的力量了!让我们从上传标签下的代码开始。我们添加了一个按钮来初始化向量索引,并将其存储在全局streamlit状态字典中,同时重置当前提取的术语。然后,从输入文本中提取术语后,我们再次将其存储在全局状态中,并给用户一个机会在插入之前审查它们。如果按下插入按钮,则调用我们的插入术语函数,更新我们插入术语的全局跟踪,并从会话状态中删除最近提取的术语。

查询提取的术语/定义

提取并保存术语和定义后,我们如何使用它们?用户甚至如何记住之前保存的内容?我们可以简单地在应用程序中添加更多标签来处理这些功能。

...
setup_tab, terms_tab, upload_tab, query_tab = st.tabs(
    ["Setup", "All Terms", "Upload/Extract Terms", "Query Terms"]
)
...
with terms_tab:
    with terms_tab:
        st.subheader("Current Extracted Terms and Definitions")
        st.json(st.session_state["all_terms"])
...
with query_tab:
    st.subheader("Query for Terms/Definitions!")
    st.markdown(
        (
            "The LLM will attempt to answer your query, and augment it's answers using the terms/definitions you've inserted. "
            "If a term is not in the index, it will answer using it's internal knowledge."
        )
    )
    if st.button("Initialize Index and Reset Terms", key="init_index_2"):
        st.session_state["llama_index"] = initialize_index(
            llm_name, model_temperature, api_key
        )
        st.session_state["all_terms"] = {}

    if "llama_index" in st.session_state:
        query_text = st.text_input("Ask about a term or definition:")
        if query_text:
            query_text = (
                query_text
                + "\nIf you can't find the answer, answer the query with the best of your knowledge."
            )
            with st.spinner("Generating answer..."):
                response = (
                    st.session_state["llama_index"]
                    .as_query_engine(
                        similarity_top_k=5,
                        response_mode="compact",
                        text_qa_template=TEXT_QA_TEMPLATE,
                        refine_template=DEFAULT_REFINE_PROMPT,
                    )
                    .query(query_text)
                )
            st.markdown(str(response))

虽然这主要是基本的,但有一些重要的事情需要注意:

我们的初始化按钮与另一个按钮有相同的文本。Streamlit会对此抱怨,所以我们提供了一个唯一的key。
一些额外的文本已添加到查询中!这是为了尝试补偿索引没有答案的时候。
在我们的索引查询中,我们指定了两个选项:
similarity_top_k=5 意味着索引将获取与查询最接近的5个术语/定义。
response_mode="compact" 意味着尽可能多的文本将从5个匹配的术语/定义中用于每次LLM调用。没有这个,索引将至少进行5次LLM调用,这可能会减慢用户的速度。

干运行测试

好吧,我希望你一直在测试。但现在,让我们尝试一次完整的测试。

刷新应用程序
输入你的LLM设置
转到查询标签
询问以下内容:什么是bunnyhug?
应用程序应该给出一些无意义的响应。如果你不知道,bunnyhug是hoodie的另一个词,由加拿大草原上的人使用!
让我们将这个定义添加到应用程序中。打开上传标签并输入以下文本:A bunnyhug是描述hoodie的常用术语。这个术语由加拿大草原上的人使用。
点击提取按钮。几分钟后,应用程序应该显示正确提取的术语/定义。点击插入术语按钮保存它!
如果我们打开术语标签,我们刚刚提取的术语和定义应该显示
返回查询标签并尝试询问什么是bunnyhug。现在,答案应该是正确的!

改进 #1 - 创建一个起始索引

有了我们的基本应用程序工作,可能会觉得建立一个有用的索引需要很多工作。如果我们给用户一些起点来展示应用程序的查询能力呢?我们可以做到这一点!首先,让我们对应用程序进行一个小改动,以便在每次上传后将索引保存到磁盘:

def insert_terms(terms_to_definition):
    for term, definition in terms_to_definition.items():
        doc = Document(text=f"Term: {term}\nDefinition: {definition}")
        st.session_state["llama_index"].insert(doc)
    # TEMPORARY - save to disk
    st.session_state["llama_index"].storage_context.persist()

现在,我们需要一些文档来提取!本项目的仓库使用了纽约市维基百科页面,你可以在这里找到文本。

如果你将文本粘贴到上传标签并运行它(可能需要一些时间),我们可以插入提取的术语。确保还将提取的术语文本复制到记事本或类似工具中!我们稍后会需要它们。

插入后,删除我们用来将索引保存到磁盘的代码行。有了现在保存的起始索引,我们可以修改我们的initialize_index函数,如下所示:

@st.cache_resource
def initialize_index(llm_name, model_temperature, api_key):
    """Load the Index object."""
    Settings.llm = get_llm(llm_name, model_temperature, api_key)

    index = load_index_from_storage(storage_context)

    return index

你还记得将那个巨大的提取术语列表保存在记事本中吗?现在当我们初始化应用程序时,我们希望将索引中的默认术语传递给我们的全局术语状态:

...
if "all_terms" not in st.session_state:
    st.session_state["all_terms"] = DEFAULT_TERMS
...

在任何我们之前重置all_terms值的地方重复上述操作。

改进 #2 - (提炼)更好的提示

如果你现在玩一下应用程序,你可能会注意到它停止遵循我们的提示了!记得我们添加到query_str变量中,如果术语/定义找不到,就用最好的知识回答。但现在如果你尝试询问随机的术语(比如bunnyhug!),它可能会也可能不会遵循这些指令。

这是由于Llama Index中“提炼”答案的概念。由于我们在前5个匹配结果中进行查询,有时所有结果不能放入一个提示中!OpenAI模型通常有4097个标记的最大输入大小。因此,Llama Index通过将匹配结果分成适合提示的块来解决这个问题。在Llama Index从第一次API调用获得初始答案后,它将下一个块发送到API,连同之前的答案,并要求模型提炼该答案。

所以,提炼过程似乎在干扰我们的结果!与其将额外指令附加到query_str,删除那个,Llama Index将让我们提供我们自己的自定义提示!现在让我们创建这些,使用默认提示和聊天特定提示作为指南。使用一个新的文件constants.py,让我们创建一些新的查询模板:

from llama_index.core import (
    PromptTemplate,
    SelectorPromptTemplate,
    ChatPromptTemplate,
)
from llama_index.core.prompts.utils import is_chat_model
from llama_index.core.llms import ChatMessage, MessageRole

# Text QA templates
DEFAULT_TEXT_QA_PROMPT_TMPL = (
    "Context information is below. \n"
    "---------------------\n"
    "{context_str}"
    "\n---------------------\n"
    "Given the context information answer the following question "
    "(if you don't know the answer, use the best of your knowledge): {query_str}\n"
)
TEXT_QA_TEMPLATE = PromptTemplate(DEFAULT_TEXT_QA_PROMPT_TMPL)

# Refine templates
DEFAULT_REFINE_PROMPT_TMPL = (
    "The original question is as follows: {query_str}\n"
    "We have provided an existing answer: {existing_answer}\n"
    "We have the opportunity to refine the existing answer "
    "(only if needed) with some more context below.\n"
    "------------\n"
    "{context_msg}\n"
    "------------\n"
    "Given the new context and using the best of your knowledge, improve the existing answer. "
    "If you can't improve the existing answer, just repeat it again."
)
DEFAULT_REFINE_PROMPT = PromptTemplate(DEFAULT_REFINE_PROMPT_TMPL)

CHAT_REFINE_PROMPT_TMPL_MSGS = [
    ChatMessage(content="{query_str}", role=MessageRole.USER),
    ChatMessage(content="{existing_answer}", role=MessageRole.ASSISTANT),
    ChatMessage(
        content="We have the opportunity to refine the above answer "
        "(only if needed) with some more context below.\n"
        "------------\n"
        "{context_msg}\n"
        "------------\n"
        "Given the new context and using the best of your knowledge, improve the existing answer. "
        "If you can't improve the existing answer, just repeat it again.",
        role=MessageRole.USER,
    ),
]

CHAT_REFINE_PROMPT = ChatPromptTemplate(CHAT_REFINE_PROMPT_TMPL_MSGS)

# refine prompt selector
REFINE_TEMPLATE = SelectorPromptTemplate(
    default_template=DEFAULT_REFINE_PROMPT,
    conditionals=[(is_chat_model, CHAT_REFINE_PROMPT)],
)

这看起来像很多代码,但并不太糟糕!如果你看了默认提示,你可能注意到有默认提示和特定于聊天的提示。继续这个趋势,我们对我们的自定义提示做同样的事情。然后,使用提示选择器,我们可以将两个提示组合成一个对象。如果使用的LLM是聊天模型(ChatGPT,GPT-4),则使用聊天提示。否则,使用正常的提示模板。

另一个需要注意的事情是,我们只定义了一个QA模板。在聊天模型中,这将转换为一个“人类”消息。

所以,现在我们可以将这些提示导入我们的应用程序并在查询期间使用它们。

from constants import REFINE_TEMPLATE, TEXT_QA_TEMPLATE

...
if "llama_index" in st.session_state:
    query_text = st.text_input("Ask about a term or definition:")
    if query_text:
        query_text = query_text  # Notice we removed the old instructions
        with st.spinner("Generating answer..."):
            response = (
                st.session_state["llama_index"]
                .as_query_engine(
                    similarity_top_k=5,
                    response_mode="compact",
                    text_qa_template=TEXT_QA_TEMPLATE,
                    refine_template=DEFAULT_REFINE_PROMPT,
                )
                .query(query_text)
            )
        st.markdown(str(response))
...

如果你更多地实验查询,希望你注意到响应现在更好地遵循我们的指令!

改进 #3 - 图像支持

Llama index还支持图像!使用Llama Index,我们可以上传文档的图像(论文、信件等),Llama Index处理提取文本。我们可以利用这一点,也允许用户上传文档的图像并从中提取术语和定义。

如果你在导入PIL时遇到错误,首先使用pip install Pillow安装它。

from PIL import Image
from llama_index.readers.file import ImageReader

@st.cache_resource
def get_file_extractor():
    image_parser = ImageReader(keep_image=True, parse_text=True)
    file_extractor = {
        ".jpg": image_parser,
        ".png": image_parser,
        ".jpeg": image_parser,
    }
    return file_extractor

file_extractor = get_file_extractor()
...
with upload_tab:
    st.subheader("Extract and Query Definitions")
    if st.button("Initialize Index and Reset Terms", key="init_index_1"):
        st.session_state["llama_index"] = initialize_index(
            llm_name, model_temperature, api_key
        )
        st.session_state["all_terms"] = DEFAULT_TERMS

    if "llama_index" in st.session_state:
        st.markdown(
            "Either upload an image/screenshot of a document, or enter the text manually."
        )
        uploaded_file = st.file_uploader(
            "Upload an image/screenshot of a document:",
            type=["png", "jpg", "jpeg"],
        )
        document_text = st.text_area("Or enter raw text")
        if st.button("Extract Terms and Definitions") and (
            uploaded_file or document_text
        ):
            st.session_state["terms"] = {}
            terms_docs = {}
            with st.spinner("Extracting (images may be slow)..."):
                if document_text:
                    terms_docs.update(
                        extract_terms(
                            [Document(text=document_text)],
                            term_extract_str,
                            llm_name,
                            model_temperature,
                            api_key,
                        )
                    )
                if uploaded_file:
                    Image.open(uploaded_file).convert("RGB").save("temp.png")
                    img_reader = SimpleDirectoryReader(
                        input_files=["temp.png"], file_extractor=file_extractor
                    )
                    img_docs = img_reader.load_data()
                    os.remove("temp.png")
                    terms_docs.update(
                        extract_terms(
                            img_docs,
                            term_extract_str,
                            llm_name,
                            model_temperature,
                            api_key,
                        )
                    )
            st.session_state["terms"].update(terms_docs)

        if "terms" in st.session_state and st.session_state["terms"]:
            st.markdown("Extracted terms")
            st.json(st.session_state["terms"])

            if st.button("Insert terms?"):
                with st.spinner("Inserting terms"):
                    insert_terms(st.session_state["terms"])
                st.session_state["all_terms"].update(st.session_state["terms"])
                st.session_state["terms"] = {}
                st.experimental_rerun()

在这里,我们使用Streamlit添加了上传文件的选项。然后打开图像并保存到磁盘(这看起来很hacky,但保持简单)。然后我们将图像路径传递给阅读器,提取文档/文本,并删除我们的临时图像文件。

现在我们有了文档,我们可以像之前一样调用extract_terms()

结论/TLDR

在本教程中,我们涵盖了大量信息,同时解决了一些常见的问题和问题:

使用不同的索引用于不同的用例(列表 vs 向量索引) 使用Streamlit的session_state概念存储全局状态值 使用Llama Index自定义内部提示 使用Llama Index从图像中读取文本

本教程的最终版本可以在这里找到,一个实时的托管演示可以在Huggingface Spaces上找到。

总结

更新时间 2024-08-23