Task-Aware Embeddings

Task-Aware Embeddings

Overview

Task-aware embeddings represent a breakthrough in semantic processing. Instead of one-size-fits-all vectors, you can optimize embeddings for specific tasks, achieving 15-30% performance improvements in search relevance, classification accuracy, and similarity matching.

This advanced feature is available across all embedding providers in Esperanto, with varying levels of native support and emulation.

The Problem with Generic Embeddings

Traditional embedding approaches use the same model configuration for all use cases:

# Traditional approach - same embedding for everything
generic_model = AIFactory.create_embedding("openai", "text-embedding-3-small")

# Query: "fast cars"
query_embedding = generic_model.embed(["fast cars"])

# Document: "high-performance vehicles"
doc_embedding = generic_model.embed(["high-performance vehicles"])

# Search result might miss the connection! 😞

The issue is that queries and documents have different semantic patterns. Queries are typically short and intent-focused, while documents are longer and information-rich. Using the same embedding strategy for both results in suboptimal matching.

The Task-Aware Solution

Task-aware embeddings optimize the model for specific purposes:

from esperanto.factory import AIFactory
from esperanto.common_types.task_type import EmbeddingTaskType

# Specialized models for different purposes
query_model = AIFactory.create_embedding(
    "jina", "jina-embeddings-v3",
    config={"task_type": EmbeddingTaskType.RETRIEVAL_QUERY}
)

document_model = AIFactory.create_embedding(
    "jina", "jina-embeddings-v3",
    config={"task_type": EmbeddingTaskType.RETRIEVAL_DOCUMENT}
)

# Now they're optimized to find each other! ✨
query_embedding = query_model.embed(["fast cars"])
doc_embedding = document_model.embed(["high-performance vehicles"])
# Much better similarity scores!

Universal Task Types

All providers support these task types through a unified interface:

Retrieval Tasks

Optimize for search and information retrieval:

# For search queries (what users type)
query_model = AIFactory.create_embedding(
    provider="any",  # Works with any provider!
    model_name="any-model",
    config={"task_type": EmbeddingTaskType.RETRIEVAL_QUERY}
)

# For documents (what gets searched)
document_model = AIFactory.create_embedding(
    provider="any",
    model_name="any-model",
    config={"task_type": EmbeddingTaskType.RETRIEVAL_DOCUMENT}
)

Best for: Search engines, RAG systems, Q&A platforms, documentation search

Performance gain: 15-25% improvement in search relevance

Classification Tasks

Optimize for categorizing and labeling text:

classification_model = AIFactory.create_embedding(
    "google", "text-embedding-004",
    config={"task_type": EmbeddingTaskType.CLASSIFICATION}
)

# Optimized for categorizing text
emails = [
    "I want to return this broken item",  # → Support
    "When will my order arrive?",         # → Shipping
    "I love this product!",               # → Feedback
]
embeddings = classification_model.embed(emails)

Best for: Email routing, content moderation, sentiment analysis, topic modeling

Performance gain: Better separation between categories, cleaner decision boundaries

Similarity & Clustering

Optimize for finding similar content:

similarity_model = AIFactory.create_embedding(
    "transformers", "all-mpnet-base-v2",
    config={"task_type": EmbeddingTaskType.SIMILARITY}
)

# Find similar content
articles = ["AI in healthcare", "Medical AI applications", "Sports news"]
embeddings = similarity_model.embed(articles)
# First two will be much closer than the third

Best for: Recommendation systems, duplicate detection, content clustering, plagiarism detection

Code Retrieval

Optimize for programming-related content:

code_model = AIFactory.create_embedding(
    "jina", "jina-embeddings-v3",
    config={"task_type": EmbeddingTaskType.CODE_RETRIEVAL}
)

# Understands programming concepts
code_snippets = [
    "def fibonacci(n): return n if n <= 1 else fib(n-1) + fib(n-2)",
    "function factorial(n) { return n <= 1 ? 1 : n * factorial(n-1) }",
    "class Car: def __init__(self, brand): self.brand = brand"
]
embeddings = code_model.embed(code_snippets)

Best for: Code search engines, API documentation, programming tutorials, code completion

Question Answering

Optimize for Q&A patterns:

qa_model = AIFactory.create_embedding(
    "google", "text-embedding-004",
    config={"task_type": EmbeddingTaskType.QUESTION_ANSWERING}
)

# Optimized for Q&A patterns
questions = ["What is machine learning?", "How does AI work?"]
answers = ["ML is a subset of AI...", "AI works by processing data..."]

q_embeddings = qa_model.embed(questions)
a_embeddings = qa_model.embed(answers)

Best for: FAQ systems, educational platforms, customer support, chatbots

Fact Verification

Optimize for fact-checking and verification:

fact_model = AIFactory.create_embedding(
    "google", "text-embedding-004",
    config={"task_type": EmbeddingTaskType.FACT_VERIFICATION}
)

# Optimized for fact-checking
claims = ["The Earth is round", "Water boils at 100°C at sea level"]
evidence = ["Scientific consensus...", "Physics textbooks state..."]

claim_embeddings = fact_model.embed(claims)
evidence_embeddings = fact_model.embed(evidence)

Best for: Fact-checking systems, misinformation detection, journalism tools, research validation

Provider Implementation Differences

Different providers implement task-aware embeddings in different ways:

Native Implementation

Providers like Jina and Google have native API support for task types. The task type is passed directly to their API:

# Jina native task support
model = AIFactory.create_embedding(
    "jina", "jina-embeddings-v3",
    config={"task_type": EmbeddingTaskType.RETRIEVAL_QUERY}
)
# Task type sent directly to Jina API

Prefix-Based Implementation

Providers without native support use intelligent prefixes:

# OpenAI uses smart prefixes
model = AIFactory.create_embedding(
    "openai", "text-embedding-3-small",
    config={"task_type": EmbeddingTaskType.RETRIEVAL_QUERY}
)
# Text is prefixed with: "Represent this query for retrieving relevant documents: "

Real-World RAG Pipeline

Here's a complete example of using task-aware embeddings in a RAG system:

from esperanto.factory import AIFactory
from esperanto.common_types.task_type import EmbeddingTaskType

# Step 1: Create specialized models
query_model = AIFactory.create_embedding(
    "jina", "jina-embeddings-v3",
    config={
        "task_type": EmbeddingTaskType.RETRIEVAL_QUERY,
        "output_dimensions": 512  # Faster search
    }
)

document_model = AIFactory.create_embedding(
    "jina", "jina-embeddings-v3",
    config={
        "task_type": EmbeddingTaskType.RETRIEVAL_DOCUMENT,
        "late_chunking": True,        # Handle long docs
        "output_dimensions": 512      # Match query model
    }
)

# Step 2: Index your knowledge base
knowledge_base = [
    "Esperanto is a unified interface for AI models...",
    "Task-aware embeddings optimize for specific use cases...",
    "RAG systems retrieve relevant context before generation..."
]

print("🔄 Indexing knowledge base...")
doc_embeddings = document_model.embed(knowledge_base)

# Step 3: Query processing
def ask_question(question):
    # Optimize query embedding for retrieval
    query_embedding = query_model.embed([question])

    # Find most relevant documents
    similarities = calculate_similarities(query_embedding, doc_embeddings)
    best_docs = get_top_k(similarities, k=3)

    return best_docs

# Step 4: Test the system
question = "How do I optimize embeddings for search?"
relevant_docs = ask_question(question)
print(f"📚 Found {len(relevant_docs)} relevant documents")

Use Cases

Multi-Task Search Engine

Different content types benefit from different task optimizations:

class MultiTaskSearchEngine:
    def __init__(self):
        # Text documents
        self.text_model = AIFactory.create_embedding(
            "jina", "jina-embeddings-v3",
            config={
                "task_type": EmbeddingTaskType.RETRIEVAL_DOCUMENT,
                "output_dimensions": 512
            }
        )

        # Code repositories
        self.code_model = AIFactory.create_embedding(
            "jina", "jina-embeddings-v3",
            config={
                "task_type": EmbeddingTaskType.CODE_RETRIEVAL,
                "output_dimensions": 512
            }
        )

        # User queries
        self.query_model = AIFactory.create_embedding(
            "jina", "jina-embeddings-v3",
            config={
                "task_type": EmbeddingTaskType.RETRIEVAL_QUERY,
                "output_dimensions": 512
            }
        )

    def index_content(self, items):
        for item in items:
            if item["type"] == "text":
                embedding = self.text_model.embed([item["content"]]).data[0].embedding
            elif item["type"] == "code":
                embedding = self.code_model.embed([item["content"]]).data[0].embedding

            item["embedding"] = embedding

        return items

Email Classification

Optimize for categorization tasks:

# Create classification-optimized model
classifier = AIFactory.create_embedding(
    "google", "text-embedding-004",
    config={"task_type": EmbeddingTaskType.CLASSIFICATION}
)

# Pre-compute category embeddings
categories = {
    "support": "Customer support and technical issues",
    "sales": "Sales inquiries and product information",
    "billing": "Billing, invoices, and payment questions"
}

category_embeddings = {
    name: classifier.embed([desc]).data[0].embedding
    for name, desc in categories.items()
}

# Classify incoming emails
def classify_email(email_text):
    email_emb = classifier.embed([email_text]).data[0].embedding

    # Find closest category
    best_category = max(
        category_embeddings.items(),
        key=lambda x: cosine_similarity(email_emb, x[1])
    )

    return best_category[0]

Best Practices

Match Task Types for Queries and Documents

When building search systems, use matching task types:

# ✅ Good - Matched task types
query_model = AIFactory.create_embedding(
    "jina", "jina-embeddings-v3",
    config={"task_type": EmbeddingTaskType.RETRIEVAL_QUERY}
)

doc_model = AIFactory.create_embedding(
    "jina", "jina-embeddings-v3",
    config={"task_type": EmbeddingTaskType.RETRIEVAL_DOCUMENT}
)

# ❌ Bad - Mismatched task types
query_model = AIFactory.create_embedding(
    "jina", "jina-embeddings-v3",
    config={"task_type": EmbeddingTaskType.CLASSIFICATION}
)

doc_model = AIFactory.create_embedding(
    "jina", "jina-embeddings-v3",
    config={"task_type": EmbeddingTaskType.RETRIEVAL_DOCUMENT}
)

Choose the Right Provider

Select providers based on task type support:

• Jina: Best for retrieval tasks with advanced features
• Google: Best for diverse task types (8 supported)
• OpenAI: Good for general-purpose with prefix optimization
• Transformers: Best for local/privacy-sensitive applications

Combine with Other Features

Task-aware embeddings work well with other advanced features:

# Combine task awareness with dimension control and late chunking
model = AIFactory.create_embedding(
    "jina", "jina-embeddings-v3",
    config={
        "task_type": EmbeddingTaskType.RETRIEVAL_DOCUMENT,
        "late_chunking": True,           # For long documents
        "output_dimensions": 512          # For faster search
    }
)

Performance Benchmarking

Measure the impact of task-aware embeddings:

import numpy as np
from esperanto.factory import AIFactory

def benchmark_task_optimization():
    """Compare generic vs task-aware embeddings"""

    # Test data
    queries = ["machine learning tutorial", "python programming guide"]
    documents = [
        "Learn ML with Python: comprehensive tutorial for beginners",
        "Python programming: complete guide to coding in Python"
    ]

    # Generic model
    generic_model = AIFactory.create_embedding("jina", "jina-embeddings-v3")

    # Task-optimized models
    query_model = AIFactory.create_embedding(
        "jina", "jina-embeddings-v3",
        config={"task_type": EmbeddingTaskType.RETRIEVAL_QUERY}
    )

    doc_model = AIFactory.create_embedding(
        "jina", "jina-embeddings-v3",
        config={"task_type": EmbeddingTaskType.RETRIEVAL_DOCUMENT}
    )

    # Calculate similarities
    def calculate_similarities(query_embs, doc_embs):
        similarities = []
        for q_emb in query_embs:
            for d_emb in doc_embs:
                sim = np.dot(q_emb, d_emb) / (np.linalg.norm(q_emb) * np.linalg.norm(d_emb))
                similarities.append(sim)
        return similarities

    # Generic approach
    generic_q_embs = [generic_model.embed([q]).data[0].embedding for q in queries]
    generic_d_embs = [generic_model.embed([d]).data[0].embedding for d in documents]
    generic_sims = calculate_similarities(generic_q_embs, generic_d_embs)

    # Task-optimized approach
    task_q_embs = [query_model.embed([q]).data[0].embedding for q in queries]
    task_d_embs = [doc_model.embed([d]).data[0].embedding for d in documents]
    task_sims = calculate_similarities(task_q_embs, task_d_embs)

    # Results
    print("🔍 Task-Aware Embedding Benchmark:")
    print(f"Generic similarities: {[f'{s:.3f}' for s in generic_sims]}")
    print(f"Task-optimized similarities: {[f'{s:.3f}' for s in task_sims]}")

    improvement = np.mean(task_sims) / np.mean(generic_sims) - 1
    print(f"📈 Average improvement: {improvement:.1%}")

# Run benchmark
benchmark_task_optimization()

See Also

• Embedding Capabilities Guide - Overview of embedding features
• Jina Provider - Native task-aware support
• Google Provider - Native task-aware support
• Transformers Features - Local task-aware implementation
• Model Discovery - Finding compatible models