txtai 追踪

txtai 是一个集语义搜索、LLM 编排和语言模型工作流于一体的嵌入式数据库。

MLflow Tracing 为 txtai 提供了自动追踪功能。通过调用 mlflow.txtai.autolog 函数可以启用 txtai 的自动追踪，MLflow 将捕获 LLM 调用、嵌入、向量搜索的追踪，并将其记录到当前的 MLflow 实验中。

首先，安装 MLflow txtai 扩展

pip install mlflow-txtai

然后，在你的 Python 代码中启用自动日志记录

import mlflow

mlflow.txtai.autolog()

示例

简单示例

展示追踪集成最简单的例子是为一个 Textractor 管道 引入工具。

import mlflow
from txtai.pipeline import Textractor

# Enable MLflow auto-tracing for txtai
mlflow.txtai.autolog()

# Optional: Set a tracking URI and an experiment
mlflow.set_tracking_uri("https://:5000")
mlflow.set_experiment("txtai")

# Define and run a simple Textractor pipeline.
textractor = Textractor()
textractor("https://github.com/neuml/txtai")

RAG

你可以轻松追踪一个 RAG 管道。

import mlflow
from txtai import Embeddings, RAG

# Enable MLflow auto-tracing for txtai
mlflow.txtai.autolog()

wiki = Embeddings()
wiki.load(provider="huggingface-hub", container="neuml/txtai-wikipedia-slim")

# Define prompt template
template = """
Answer the following question using only the context below. Only include information
specifically discussed.

question: {question}
context: {context} """

# Create RAG pipeline
rag = RAG(
    wiki,
    "hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
    system="You are a friendly assistant. You answer questions from users.",
    template=template,
    context=10,
)

rag("Tell me about the Roman Empire", maxlength=2048)

智能体

你可以毫不费力地追踪一个旨在研究天文学问题的 txtai 智能体 的内部运行情况。

import mlflow
from txtai import Agent, Embeddings

# Enable MLflow auto-tracing for txtai
mlflow.txtai.autolog()


def search(query):
    """
    Searches a database of astronomy data.

    Make sure to call this tool only with a string input, never use JSON.

    Args:
        query: concepts to search for using similarity search

    Returns:
        list of search results with for each match
    """

    return embeddings.search(
        "SELECT id, text, distance FROM txtai WHERE similar(:query)",
        10,
        parameters={"query": query},
    )


embeddings = Embeddings()
embeddings.load(provider="huggingface-hub", container="neuml/txtai-astronomy")

agent = Agent(
    tools=[search],
    llm="hugging-quants/Meta-Llama-3.1-8B-Instruct-AWQ-INT4",
    max_iterations=10,
)

researcher = """
{command}

Do the following.
 - Search for results related to the topic.
 - Analyze the results
 - Continue querying until conclusive answers are found
 - Write a Markdown report
"""

agent(
    researcher.format(
        command="""
Write a detailed list with explanations of 10 candidate stars that could potentially be habitable to life.
"""
    ),
    maxlength=16000,
)

更多信息

有关将 txtai 与 MLflow 结合使用的更多示例和指导，请参阅 MLflow txtai 扩展文档