LangChain TitanTakeoff Embeddings Integration

LangChain provides integration with TitanML's Takeoff server through the TitanTakeoffEmbed class, allowing you to generate embeddings from various models in your local environment.

What is TitanTakeoffEmbed?

TitanTakeoffEmbed is a class that interfaces with TitanML's Takeoff Inference API to generate embeddings from text. It provides functionality to embed both individual queries and collections of documents using locally deployed embedding models.

The class supports both synchronous and asynchronous operations, making it flexible for different use cases. It also allows you to configure multiple embedding models through consumer groups.

Reference

Here are the main methods available in TitanTakeoffEmbed:

How to Use TitanTakeoffEmbed

Installation

Before using TitanTakeoffEmbed, you need to install the takeoff-client:

pip install 'takeoff-client==0.4.0'

Basic Usage

Here's how to create embeddings using the default configuration:

from langchain_community.embeddings import TitanTakeoffEmbed

# Initialize with default settings (localhost:3000)
embeddings = TitanTakeoffEmbed()

# Generate embedding for a single query
query_embedding = embeddings.embed_query(
    "What is the weather in London?",
    consumer_group="embed"
)

# Generate embeddings for multiple documents
documents = [
    "The weather in London is rainy",
    "The weather in Paris is sunny"
]
doc_embeddings = embeddings.embed_documents(
    documents,
    consumer_group="embed"
)

Starting Custom Embedding Models

You can initialize TitanTakeoffEmbed with specific models:

import time
from langchain_community.embeddings import TitanTakeoffEmbed

# Configure the embedding model
model_config = {
    "model_name": "BAAI/bge-large-en-v1.5",
    "device": "cpu",
    "consumer_group": "embed"
}

# Initialize with custom model
embeddings = TitanTakeoffEmbed(models=[model_config])

# Wait for model to load
time.sleep(60)

# Generate embeddings using the custom model
embedding = embeddings.embed_query(
    "What is the capital of France?",
    consumer_group="embed"
)

Using Async Methods

For applications that benefit from asynchronous operations:

import asyncio
from langchain_community.embeddings import TitanTakeoffEmbed

async def generate_embeddings():
    embeddings = TitanTakeoffEmbed()
    
    # Generate single embedding asynchronously
    query_embedding = await embeddings.aembed_query(
        "What is machine learning?"
    )
    
    # Generate multiple embeddings asynchronously
    texts = [
        "Machine learning is a subset of AI",
        "Deep learning uses neural networks",
        "Natural Language Processing deals with text"
    ]
    doc_embeddings = await embeddings.aembed_documents(texts)
    
    return query_embedding, doc_embeddings

# Run the async function
query_emb, doc_embs = asyncio.run(generate_embeddings())

Custom Server Configuration

If your Takeoff server runs on different ports or hosts:

embeddings = TitanTakeoffEmbed(
    base_url="http://custom-host",
    port=8000,
    mgmt_port=8001
)

Remember that TitanTakeoff needs to be running in the background before using any of these examples. The service handles the actual embedding computation, while the TitanTakeoffEmbed class provides a convenient interface to interact with it through LangChain.