Tracing txtai

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

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

要开始，请安装 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)

Agent

您可以轻松跟踪一个用于研究天文学问题的 txtai agent 的内部工作。

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 扩展文档