追踪 txtai

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

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

开始之前，请安装 MLflow txtai 扩展

pip install mlflow-txtai

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

import mlflow

mlflow.txtai.autolog()

基本示例

第一个示例追踪了一个 Textractor pipeline。

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("http://localhost:5000")
mlflow.set_experiment("txtai")

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

检索增强生成 (RAG)

下一个示例追踪了一个 RAG pipeline。

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 (智能体)，该 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 扩展文档