Aller au contenu principal
Dernière mise à jour le

Traçage txtai

Traçage txtai via log automatique

txtai est une base de données d'embeddings tout-en-un pour la recherche sémantique, l'orchestration LLM et les workflows de modèles linguistiques.

MLflow Tracing fournit une capacité de traçage automatique pour txtai. Le traçage automatique pour txtai peut être activé en appelant la fonction mlflow.autolog, MLflow capturera les traces pour l'invocation de LLM, les embeddings, la recherche AI et les enregistrera dans l'Experimentation MLflow actif.

Prérequis​

Pour utiliser MLflow Tracing avec txtai, vous devez installer MLflow, la bibliothèque txtai et l'extension mlflow-txtai.

Pour les environnements de développement, installez le package MLflow complet avec les extras Databricks, txtai et mlflow-txtai:

Bash
pip install --upgrade "mlflow[databricks]>=3.1" txtai mlflow-txtai

Le package mlflow[databricks] complet inclut toutes les fonctionnalités pour le développement local et l’expérimentation sur Databricks.

remarque

MLflow 3 est fortement recommandé pour la meilleure expérience de traçage avec txtai.

Avant d'exécuter les exemples, vous devrez configurer votre environnement :

Pour les utilisateurs en dehors des notebooks Databricks : Définissez vos variables d'environnement Databricks :

Bash
export DATABRICKS_HOST="https://your-workspace.cloud.databricks.com"
export DATABRICKS_TOKEN="your-personal-access-token"

Pour les utilisateurs dans les Notebooks Databricks : ces identifiants sont définis automatiquement pour vous.

Clés API : Assurez-vous que les clés API de votre fournisseur LLM sont définies :

Bash
export OPENAI_API_KEY="your-openai-api-key"
# Add other provider keys as needed if using txtai with different models

Exemple de base​

Le premier exemple retrace un pipeline Textractor.

remarque

Sur les clusters de compute Serverless, la journalisation automatique pour les cadres de traçage genAI n'est pas activée automatiquement. Vous devez activer explicitement l'autologging en appelant la fonction mlflow.<library>.autolog() appropriée pour les intégrations spécifiques que vous souhaitez suivre.

Python
import mlflow
from txtai.pipeline import Textractor
import os

# Ensure any necessary LLM provider API keys are set in your environment if Textractor uses one
# For example, if it internally uses OpenAI:
# os.environ["OPENAI_API_KEY"] = "your-openai-key"

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

# Set up MLflow tracking to Databricks
mlflow.set_tracking_uri("databricks")
mlflow.set_experiment("/Shared/txtai-demo")

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

Traçage txtai Textractor via autolog

Génération augmentée de récupération (RAG)​

L'exemple suivant trace un pipeline RAG.

Python
import mlflow
from txtai import Embeddings, RAG
import os

# Ensure your LLM provider API key (e.g., OPENAI_API_KEY for the Llama model via some services) is set
# os.environ["OPENAI_API_KEY"] = "your-key" # Or HUGGING_FACE_HUB_TOKEN, etc.

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

# Set up MLflow tracking to Databricks if not already configured
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Shared/txtai-rag-demo")

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)

Traçage txtai Rag via le log automatique

Agent​

Le dernier exemple exécute un agent txtai conçu pour rechercher des questions sur l'astronomie.

Python
import mlflow
from txtai import Agent, Embeddings
import os

# Ensure your LLM provider API key (e.g., OPENAI_API_KEY for the Llama model via some services) is set
# os.environ["OPENAI_API_KEY"] = "your-key" # Or HUGGING_FACE_HUB_TOKEN, etc.

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

# Set up MLflow tracking to Databricks if not already configured
# mlflow.set_tracking_uri("databricks")
# mlflow.set_experiment("/Shared/txtai-agent-demo")

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={&quot;query&quot;: 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,
)

Traçage de l&#39;Agent txtai via autolog

Ressources supplémentaires​

Pour plus d'exemples et de conseils sur l'utilisation de txtai avec MLflow, consultez la documentation sur l'extension txtai de MLflow.

Ressources supplémentaires​

Sur cette page