Embedding Models

async def embed_summaries(
    summaries: list[ConversationSummary], 
    embedding_model: BaseEmbeddingModel
) -> list[dict[str, Union[ConversationSummary, list[float]]]]:
    """Embeds conversation summaries and returns items ready for clustering."""

[
    {
        "item": ConversationSummary(...),
        "embedding": [0.123, -0.456, 0.789, ...]  # 1536 dimensions for OpenAI
    },
    ...
]

from kura.embedding import OpenAIEmbeddingModel

embedding_model = OpenAIEmbeddingModel(
    model_name="text-embedding-3-small",
    model_batch_size=50,
    n_concurrent_jobs=5
)

@retry(wait=wait_fixed(3), stop=stop_after_attempt(3))
async def _embed_batch(self, texts: list[str]) -> list[list[float]]:
    """Embed a single batch of texts."""
    async with self._semaphore:
        resp = await self.client.embeddings.create(
            input=texts, 
            model=self.model_name
        )
        embeddings = [item.embedding for item in resp.data]
        return embeddings

from kura.embedding import SentenceTransformerEmbeddingModel

embedding_model = SentenceTransformerEmbeddingModel(
    model_name="all-MiniLM-L6-v2",
    model_batch_size=128,
    device="cuda"  # or "cpu"
)

from kura.embedding import CohereEmbeddingModel

embedding_model = CohereEmbeddingModel(
    model_name="embed-v4.0",
    model_batch_size=96,
    n_concurrent_jobs=5,
    input_type="clustering",
    api_key="your-api-key"  # or set COHERE_API_KEY env var
)

from kura.embedding import embed_summaries, OpenAIEmbeddingModel

embedding_model = OpenAIEmbeddingModel()

embedded_items = await embed_summaries(
    summaries=summaries,
    embedding_model=embedding_model
)

print(embedded_items[0])
# {
#   "item": ConversationSummary(...),
#   "embedding": [0.123, -0.456, ...]
# }

embedded_items = await embed_summaries(summaries, embedding_model)
cluster_id_to_summaries = clustering_method.cluster(embedded_items)

from kura.utils import batch_texts

# Split texts into batches
batches = batch_texts(texts, batch_size=50)
# [[text1, text2, ..., text50], [text51, ..., text100], ...]

# Process batches concurrently
tasks = [self._embed_batch(batch) for batch in batches]
results = await asyncio.gather(*tasks)

async def embed(self, texts: list[str]) -> list[list[float]]:
    # Create batches
    batches = batch_texts(texts, self._model_batch_size)
    
    # Process all batches concurrently
    tasks = [self._embed_batch(batch) for batch in batches]
    results_list_of_lists = await gather(*tasks)
    
    # Flatten results
    embeddings = []
    for result_batch in results_list_of_lists:
        embeddings.extend(result_batch)
    
    return embeddings

from kura.base_classes import BaseEmbeddingModel

class CustomEmbeddingModel(BaseEmbeddingModel):
    def slug(self) -> str:
        """Unique identifier for this model configuration."""
        return f"custom-model-{self.version}"
    
    async def embed(self, texts: list[str]) -> list[list[float]]:
        """Convert texts to embeddings.
        
        Args:
            texts: List of text strings to embed
            
        Returns:
            List of embedding vectors (one per input text)
        """
        # Your implementation here
        embeddings = await self.api_client.embed_batch(texts)
        return embeddings

# OpenAI: Larger batches = fewer API calls
embedding_model = OpenAIEmbeddingModel(
    model_batch_size=2048  # Max allowed by OpenAI
)

# Sentence Transformers: Tune based on GPU memory
embedding_model = SentenceTransformerEmbeddingModel(
    model_batch_size=256,  # Increase if you have GPU memory
    device="cuda"
)

# OpenAI: Balance rate limits vs speed
embedding_model = OpenAIEmbeddingModel(
    n_concurrent_jobs=10  # Higher = faster, but may hit rate limits
)

# Monitor rate limits: 5,000 RPM typical for most accounts

# Embeddings are automatically saved when using checkpoint managers
from kura.checkpoints import HFDatasetCheckpointManager

checkpoint_mgr = HFDatasetCheckpointManager("./checkpoints")

# First run: computes embeddings
summaries = await summarise_conversations(
    conversations=conversations,
    model=summary_model,
    checkpoint_manager=checkpoint_mgr
)

# Embeddings are in summary.metadata["embedding"] if cached

import logging

logging.basicConfig(level=logging.INFO)

# Output:
# INFO:kura.embedding:Initialized OpenAIEmbeddingModel with model=text-embedding-3-small
# INFO:kura.embedding:Starting embedding of 1000 texts using text-embedding-3-small
# DEBUG:kura.embedding:Split 1000 texts into 20 batches of size 50
# INFO:kura.embedding:Successfully embedded 1000 texts, produced 1000 embeddings