"""
Processors for different models
"""
import functools
import json
import logging
from typing import List

import cohere as co
import httpx
import openai
import tiktoken

from emb3d import config
from emb3d.types import Backend, EmbedJob, EmbedResponse, Failure, Result, WaitFor

HF_HEADERS = {}
OPENAI_INIT_PARAMS = {}
COHERE_INIT_PARAMS = {}

OPENAI_ENDPOINT = "https://api.openai.com/v1/embeddings"

_cleanup_callables = []


@functools.cache
def httpx_client() -> httpx.AsyncClient:
    """Cached httpx client"""
    cli = httpx.AsyncClient()
    _cleanup_callables.append(cli.aclose)
    return cli


@functools.cache
def cohere_client(api_key: str) -> co.AsyncClient:
    """Cached cohere client"""
    cli = co.AsyncClient(api_key=api_key)
    _cleanup_callables.append(cli.close)
    return cli


def hf_headers(job: EmbedJob) -> dict:
    """Returns the headers for the HuggingFace API"""
    return {
        "Authorization": f"Bearer {job.api_key}",
    }


async def cleanup():
    """Cleanup any resources used by the clients"""
    for cleanup_fn in _cleanup_callables:
        await cleanup_fn()


def hf_inference_url(model_id: str):
    """Inference URL for the HuggingFace API"""
    return (
        f"https://api-inference.huggingface.co/pipeline/feature-extraction/{model_id}"
    )


async def _huggingface(job: EmbedJob, inputs: List[str]) -> EmbedResponse:
    data = {"inputs": inputs, "wait_for_model": True}
    response = await httpx_client().post(
        hf_inference_url(job.model_id), headers=hf_headers(job), json=data
    )

    logging.debug("Model: %s, Response: %s", job.model_id, response.status_code)

    try:
        json_response = response.json()
    except json.JSONDecodeError:
        return Failure("HF response is not in JSON format, {response.content}")
    if response.status_code == 200:
        return Result(json_response)
    elif response.status_code == 503:
        estimated_time = json_response.get("estimated_time")
        if estimated_time:
            return WaitFor(estimated_time, None)
        else:
            return Failure(
                "[HuggingFace] Service unavailable and estimated time not provided"
            )
    return Failure(f"[HuggingFace] Unexpected response: {json_response}")


async def _openai(job: EmbedJob, inputs: List[str]) -> EmbedResponse:
    try:
        resp = await openai.Embedding.acreate(model=job.model_id, input=inputs)
        return Result([row.embedding for row in resp.data])
    except openai.error.RateLimitError as err:
        logging.debug("[OpenAI] Rate limit error: %s", err)
        return WaitFor(config.RATE_LIMIT_WAIT_TIME_SECS, error=err)
    except openai.error.APIError as err:
        logging.debug("[OpenAI] Error:", err)
        return Failure(f"[OpenAI] Error: {err}")


async def _cohere(job: EmbedJob, inputs: List[str]) -> EmbedResponse:
    cli = cohere_client(job.api_key)
    try:
        co_resp = await cli.embed(inputs, job.model_id)
        return Result(co_resp.embeddings)
    except co.error.CohereError as err:
        return Failure(f"[Cohere] Error: {err}")


async def gen(job: EmbedJob, inputs: List[str]) -> EmbedResponse:
    """
    Generate embeddings for the given inputs.

    Routes to the appropriate client based on the job specification.
    """
    if job.backend == Backend.HUGGINGFACE:
        return await _huggingface(job, inputs)
    elif job.backend == Backend.OPENAI:
        openai.api_key = job.api_key
        return await _openai(job, inputs)
    elif job.backend == Backend.COHERE:
        return await _cohere(job, inputs)
    else:
        raise ValueError(f"Unknown backend: {job.backend}")


@functools.cache
def get_encoder(model_id: str) -> tiktoken.Encoding:
    """Returns encoder used for the given model_id"""
    return tiktoken.encoding_for_model(model_id)


def approx_token_count(job: EmbedJob, input: str) -> int:
    """
    Helper to compute the approximate token count for an input

    For backends other than OpenAI, we use an approximate method as a fast
    tokenizer isn't always available.
    """
    if job.backend == Backend.OPENAI:
        encoder = get_encoder(job.model_id)
        return len(encoder.encode(input))
    return len(input) // 2

emb3d/client.py
akhilravidas/emb3d
"""
Job Execution
"""
import asyncio
import threading

from rich import print
from rich.console import Console
from rich.live import Live

from emb3d import textui
from emb3d.compute import local, remote
from emb3d.types import EmbedJob


def execute(job: EmbedJob):
    console = Console()
    console.rule("Starting Job")
    with Live(auto_refresh=True, console=console) as live:
        print("Job Config:")
        print(job.describe())
        if job.execution_config.is_remote:
            asyncio.run(remote.run(job, textui.render_ui_async(job, live)))
        else:
            ui_thread = threading.Thread(target=textui.render_ui_sync, args=(job, live))
            ui_thread.start()
            local.run(job)
            ui_thread.join()
    console.rule("Job Complete")

emb3d/compute/__init__.py
akhilravidas/emb3d
import json
import logging
from typing import Iterator

from emb3d import client
from emb3d.io import reader
from emb3d.types import Batch, EmbedJob


def write_batch_results_post_lock(job: EmbedJob, batch: Batch):
    """
    Write the results of a batch to the output file, assumes calling context has
    ensured that there is atmost one writer writing to the output file.
    """
    logging.debug("Writing computed batch results, size = [%d]", len(batch.row_ids))
    for idx, _ in enumerate(batch.row_ids):
        job.out_file.write(
            json.dumps(
                {
                    "row_id": batch.row_ids[idx],
                    "input": batch.inputs[idx],
                    "embedding": batch.embeddings[idx]
                    if batch.embeddings is not None
                    else None,
                    "error": str(batch.error) if batch.error else None,
                }
            )
            + "\n"
        )
    job.batch_saved(len(batch.row_ids))


def gen_batch(job: EmbedJob, batch_size: int, max_tokens: int) -> Iterator[Batch]:
    """
    Generates batches of rows from the input file.

    Batches are constructed so that:
    - Each batch contains atmost `batch_size` rows.
    - Each batch have atmost max_tokens (except when a single line exceeds token limit)
    """
    batch_ids = []
    batch_inputs = []
    batch_token_count = 0
    for line_num, record in enumerate(reader.jsonl(job.in_file)):
        text = record[job.column_name]
        new_tokens = client.approx_token_count(job, text)

        can_merge_token = batch_token_count + new_tokens < max_tokens
        can_merge_element = len(batch_ids) + 1 < batch_size

        can_merge = len(batch_ids) == 0 or (can_merge_token and can_merge_element)
        if can_merge:
            batch_ids.append(line_num)
            batch_inputs.append(text)
            batch_token_count += new_tokens
        else:
            yield Batch(batch_ids, batch_inputs)
            batch_ids = [line_num]
            batch_inputs = [text]
            batch_token_count = new_tokens

    if batch_ids:
        yield Batch(batch_ids, batch_inputs)

emb3d/compute/common.py
akhilravidas/emb3d
import sentence_transformers

from emb3d import config
from emb3d.compute.common import gen_batch, write_batch_results_post_lock
from emb3d.types import EmbedJob


def run(job: EmbedJob):
    """
    Run the job.
    """
    model = sentence_transformers.SentenceTransformer(job.model_id)
    for batch in gen_batch(job, job.batch_size, config.max_tokens(job.backend)):
        batch.embeddings = model.encode(batch.inputs).tolist()
        batch.error = None
        write_batch_results_post_lock(job, batch)

emb3d/compute/local.py
akhilravidas/emb3d
import asyncio
import logging
from typing import Coroutine

from aiolimiter import AsyncLimiter

from emb3d import client, config, textui
from emb3d.compute.common import gen_batch, write_batch_results_post_lock
from emb3d.types import Batch, EmbedJob, Failure, Result, WaitFor


async def terminate(*tasks):
    """Cancels tasks and waits for them to complete."""
    logging.debug("Terminating tasks...")
    for task in tasks:
        task.cancel()
    for task in tasks:
        try:
            await task
        except asyncio.CancelledError:
            pass


async def produce(job: EmbedJob, queue: asyncio.Queue):
    """
    Producer task that generates batches and pushes them to the queue.
    """
    logging.debug("Producer: Starting")
    for batch in gen_batch(job, job.batch_size, config.max_tokens(job.backend)):
        logging.debug("Producer: Next Batch [%d]", len(batch.row_ids))
        await queue.put(batch)


FILE_WRITE_LOCK = asyncio.Lock()


async def write_batch_results(job: EmbedJob, batch: Batch):
    """Write batch results to output file."""
    async with FILE_WRITE_LOCK:
        # We use synchronous file operations as the data being written is small
        write_batch_results_post_lock(job, batch)


async def worker(
    job: EmbedJob, rate_limiter: AsyncLimiter, job_queue: asyncio.Queue, num_retries=2
):
    """
    Consumer task that consumes batches from the queue and generates embeddings.
    """
    while True:
        batch = await job_queue.get()
        assert len(batch.inputs) == len(batch.row_ids)
        batch_retry = num_retries
        while batch_retry > 0:
            batch_retry -= 1
            await rate_limiter.acquire()
            job.tracker.encoding += len(batch.inputs)
            resp = await client.gen(job, batch.inputs)
            job.tracker.encoding -= len(batch.inputs)
            # match resp:
            #     case Result(data):
            #         assert len(data) == len(batch.inputs)
            #         # clear off transient error
            #         batch.error = None
            #         batch.embeddings = data
            #         job.batch_success(len(batch.inputs))
            #         await write_batch_results(job, batch)
            #         break
            #     case Failure(error):
            #         batch.error = error
            #     case WaitFor(seconds, error):
            #         batch.error = error
            #         if batch_retry != 0:
            #             await asyncio.sleep(seconds)
            if isinstance(resp, Result):
                assert len(resp.data) == len(batch.inputs)
                # clear off transient error
                batch.error = None
                batch.embeddings = resp.data
                job.batch_success(len(batch.inputs))
                await write_batch_results(job, batch)
                break
            elif isinstance(resp, Failure):
                batch.error = resp.error
            elif isinstance(resp, WaitFor):
                batch.error = resp.error
                if batch_retry != 0:
                    await asyncio.sleep(resp.seconds)

            if batch_retry == 0:
                job.batch_failure(len(batch.inputs))
                if batch.error is not None:
                    job.batch_error(str(batch.error))
                await write_batch_results(job, batch)
        job_queue.task_done()


async def consume(
    job: EmbedJob, rate_limiter: AsyncLimiter, job_queue: asyncio.Queue, num_retries=2
):
    """
    Consumer task that consumes batches from the queue and generates embeddings.
    """
    logging.debug("Starting consumer task")
    return await asyncio.gather(
        *[
            worker(job, rate_limiter, job_queue, num_retries)
            for _ in range(job.max_concurrent_requests)
        ],
        return_exceptions=True,
    )


async def run(job: EmbedJob, update_ui_coroutine: Coroutine[None, None, None]):
    """
    Main entry point for the embedding job.
    """
    job_queue = asyncio.Queue(maxsize=job.max_concurrent_requests)
    ui_task = asyncio.create_task(update_ui_coroutine)
    request_limiter = AsyncLimiter(config.max_requests_per_minute(job.backend), 60)
    producer_task = asyncio.create_task(produce(job, job_queue))
    consumer_task = asyncio.create_task(
        consume(job, request_limiter, job_queue=job_queue)
    )
    try:
        await asyncio.wait({producer_task}, return_when=asyncio.FIRST_EXCEPTION)
        await job_queue.join()
        await terminate(consumer_task)
    except KeyboardInterrupt:
        await terminate(producer_task, consumer_task)
    finally:
        await client.cleanup()
        await ui_task

emb3d/compute/remote.py
akhilravidas/emb3d
import io
import json

from emb3d.compute.common import gen_batch, write_batch_results_post_lock
from emb3d.test_utils import mock_embed_job
from emb3d.types import Batch


def test_write_batch_results_post_lock():
    job = mock_embed_job()
    batch = Batch(row_ids=[1], inputs=["hello"], embeddings=[[4, 2]], error=None)

    write_batch_results_post_lock(job, batch)

    job.out_file.seek(0)
    written_data = json.loads(job.out_file.read())

    assert written_data["row_id"] == 1
    assert written_data["input"] == "hello"
    assert written_data["embedding"] == [4, 2]
    assert written_data["error"] is None


def test_gen_batch():
    in_file = io.StringIO('{"text": "hello"}\n{"text": "world"}')
    job = mock_embed_job(in_file=in_file)

    batches = list(gen_batch(job, batch_size=1, max_tokens=5))

    assert len(batches) == 2
    assert batches[0].row_ids == [0]
    assert batches[0].inputs == ["hello"]
    assert batches[1].row_ids == [1]
    assert batches[1].inputs == ["world"]


def test_gen_batch_low_tokens():
    in_file = io.StringIO('{"text": "hello"}\n{"text": "world"}')
    job = mock_embed_job(in_file=in_file)

    batches = list(gen_batch(job, batch_size=1, max_tokens=1))

    # Each batch should have atleast one element when if its above the token limit
    assert len(batches) == 2
    assert batches[0].row_ids == [0]
    assert batches[0].inputs == ["hello"]
    assert batches[1].row_ids == [1]
    assert batches[1].inputs == ["world"]

    in_file.seek(0)
    batches = list(gen_batch(job, batch_size=10, max_tokens=1))

    # Each batch should have atleast one element when if its above the token limit
    assert len(batches) == 2
    assert batches[0].row_ids == [0]
    assert batches[0].inputs == ["hello"]
    assert batches[1].row_ids == [1]
    assert batches[1].inputs == ["world"]


def test_gen_batch_high_tokens():
    in_file = io.StringIO('{"text": "hello"}\n{"text": "world"}')
    job = mock_embed_job(in_file=in_file)

    batches = list(gen_batch(job, batch_size=1, max_tokens=100))

    assert len(batches) == 2
    assert batches[0].row_ids == [0]
    assert batches[0].inputs == ["hello"]
    assert batches[1].row_ids == [1]
    assert batches[1].inputs == ["world"]

    in_file.seek(0)
    batches = list(gen_batch(job, batch_size=10, max_tokens=10))

    assert len(batches) == 1
    assert batches[0].row_ids == [0, 1]
    assert batches[0].inputs == ["hello", "world"]

emb3d/compute/test_common.py
akhilravidas/emb3d
import io
import json
from io import reader
from unittest.mock import Mock, patch

import numpy as np
import pytest

from emb3d.compute.local import run
from emb3d.test_utils import mock_embed_job


@pytest.mark.parametrize(
    "inputs, expected_embeddings",
    [
        (["Hello world!", "Bye world!"], np.array([[0.1, 0.2], [0.5, 0.6]])),
        (["Testing 123"], np.array([[0.7, 0.8]])),
    ],
)
def test_run(inputs, expected_embeddings):
    mock_model = Mock()
    mock_model.encode.return_value = expected_embeddings

    with patch(
        "sentence_transformers.SentenceTransformer", return_value=mock_model
    ) as mock_transformer:
        records = [{"text": input} for input in inputs]
        in_file = io.StringIO("\n".join(json.dumps(record) for record in records))
        job = mock_embed_job(in_file=in_file, batch_size=100)
        run(job)

        mock_transformer.assert_called_once_with(job.model_id)
        mock_model.encode.assert_called_with(inputs)
        out_file = job.out_file
        out_file.seek(0)
        saved_records = list(reader.jsonl(out_file))
        print(saved_records)
        assert len(saved_records) == len(inputs)
        for idx, _ in enumerate(inputs):
            assert saved_records[idx]["embedding"] == expected_embeddings[idx].tolist()

emb3d/compute/test_local.py
akhilravidas/emb3d
from pathlib import Path
from typing import Optional

import altair as alt
import hdbscan
import numpy as np
import pandas as pd
from umap import UMAP

from emb3d.io import reader

NUM_TITLES = 20
READ_CHUNK_SIZE = 500


def cluster_hdbscan(X: pd.DataFrame, min_cluster_size: int) -> hdbscan.HDBSCAN:
    return hdbscan.HDBSCAN(min_cluster_size=min_cluster_size).fit(X)


# TODO: Very very inefficient, time and heap allocation wise
def get_data(embedding_file: Path, label_field: Optional[str]):
    embeddings = []
    titles = []

    chunk_iter = pd.read_json(embedding_file, lines=True, chunksize=READ_CHUNK_SIZE)

    offset = 0
    for chunk in chunk_iter:
        embeddings.extend(chunk["embedding"].tolist())
        titles.extend(
            chunk.get(label_field, chunk.get("id", chunk.index + offset)).tolist()
        )
        offset += len(chunk)

    df_embeddings = pd.DataFrame(embeddings)

    return df_embeddings, titles


def umap_reduce(X: pd.DataFrame) -> np.ndarray:
    reducer = UMAP()
    return reducer.fit_transform(X)  # type: ignore


def generate_chart(X_reduced, titles, hdbscan_model) -> alt.TopLevelMixin:
    if hdbscan_model:
        cluster_labels = hdbscan_model.labels_
        df_cluster_titles = pd.DataFrame(
            {
                "x1": X_reduced[:, 0],
                "x2": X_reduced[:, 1],
                "title": titles,
                "cluster": cluster_labels,
            }
        )

        df_agg = (
            df_cluster_titles.groupby("cluster")
            .agg(
                x1=("x1", "mean"),
                x2=("x2", "mean"),
                count=("cluster", "size"),
                cluster_title=("title", "first"),
            )
            .reset_index()
        )

        data = df_agg
    else:
        data = {
            "x1": X_reduced[:, 0],
            "x2": X_reduced[:, 1],
            "title": titles,
        }
        data = pd.DataFrame(data)

    brush = alt.selection_interval()

    scatter_plot = (
        alt.Chart(data)
        .mark_circle(opacity=0.6)
        .encode(
            x=alt.X("x1", axis=None, scale=alt.Scale(zero=False)),
            y=alt.Y("x2", axis=None, scale=alt.Scale(zero=False)),
            tooltip=["cluster_title", "count"] if hdbscan_model else ["title"],
            color=alt.Color("cluster:N", legend=None)
            if hdbscan_model
            else alt.value("blue"),
            size=alt.Size("count:Q", legend=None, scale=alt.Scale(range=[10, 200]))
            if hdbscan_model
            else alt.value(20),
        )
        .properties(width=1000)
        .add_params(brush)
    )

    title_list = (
        alt.Chart(data)
        .mark_text(align="left")
        .encode(
            y=alt.Y("row_number:O", axis=None),
            text="cluster_title:N" if hdbscan_model else "title:N",
        )
        .transform_window(row_number="row_number()")
        .transform_filter(brush)
        .transform_window(rank="rank(row_number)")
        .transform_filter(alt.datum.rank < NUM_TITLES)
    )

    return (
        alt.vconcat(scatter_plot, title_list)
        .resolve_legend(color="independent")
        .configure_view(stroke=None)
    )

emb3d/compute/visualize.py
akhilravidas/emb3d
"""
App and Run Configuration parameters
"""
from __future__ import annotations

import functools
import os
from dataclasses import dataclass, fields
from pathlib import Path
from typing import Optional

import yaml

from emb3d.types import Backend

# Scale down factor
SCALE_DOWN_FACTOR = 0.8

VISUALIZATION_CLUSTERING_THRESHOLD = 5000
VISUALIZATION_DEFAULT_MIN_CLUSTER_SIZE = 10

max_requests_limits = {
    Backend.OPENAI: 10000,
    Backend.COHERE: 1500,  # Using trial key values
    Backend.HUGGINGFACE: 100,
}

RATE_LIMIT_WAIT_TIME_SECS = 0.5

max_token_limits = {
    Backend.OPENAI: 8191,
    Backend.COHERE: 8000,
    Backend.HUGGINGFACE: 512,
}


@functools.cache
def max_requests_per_minute(backend: Backend) -> int:
    base_rpm = max_requests_limits.get(backend, 100)
    return int(base_rpm * SCALE_DOWN_FACTOR)


def max_tokens(backend: Backend) -> int:
    return max_token_limits.get(backend, 512)


def app_data_root() -> Path:
    sys_data_root = Path.home() / ".cache"
    return Path(os.getenv("XDG_CACHE_HOME", sys_data_root)) / "emb3d"


def config_path() -> Path:
    return app_data_root() / "config.yaml"


@dataclass
class AppConfig:
    default_model: str = "text-embedding-ada-002"
    openai_token: Optional[str] = None
    cohere_token: Optional[str] = None
    huggingface_token: Optional[str] = None

    @classmethod
    def instance(cls) -> AppConfig:
        if hasattr(cls, "_config"):
            return cls._config

        if not config_path().exists():
            cls._config = cls()
            return cls._config
        else:
            with config_path().open() as f:
                cls._config = cls(**yaml.safe_load(f))
                return cls._config

    def save(self):
        if not os.path.exists(app_data_root()):
            os.makedirs(app_data_root())
        with config_path().open("w") as f:
            yaml.safe_dump(self.__dict__, f)

    def get(self, key: str) -> Optional[str]:
        valid_keys = {field.name for field in fields(AppConfig)}
        if key not in valid_keys:
            raise ValueError(f"'{key}' is not a valid configuration option")

        return getattr(self, key, None)

    def set(self, key: str, value: str):
        valid_keys = {field.name for field in fields(AppConfig)}
        if key not in valid_keys:
            raise ValueError(f"'{key}' is not a valid configuration option")

        setattr(self, key, value)
        self.save()

emb3d/config.py
akhilravidas/emb3d
"""
Readers
"""
import json
from typing import Iterator, TextIO


def line(f_io: TextIO) -> Iterator[str]:
    for line in f_io:
        yield line.strip()


def jsonl(f_io: TextIO) -> Iterator[dict]:
    for nxt_line in line(f_io):
        yield json.loads(nxt_line)

emb3d/io/reader.py
akhilravidas/emb3d
"""
Writers
"""
from pathlib import Path

import altair as alt


def chart2html(chart: alt.TopLevelMixin, out_fname: Path):
    """
    Write chart to html file
    """
    with alt.data_transformers.enable("default"):
        chart.save(out_fname)

emb3d/io/writer.py
akhilravidas/emb3d
"""
emb3d CLI
"""
import os
import random
import string
import sys
import webbrowser
from enum import Enum
from io import StringIO
from pathlib import Path
from typing import Optional, TextIO

import typer
from rich.prompt import Prompt
from typing_extensions import Annotated

from emb3d import compute, config, textui
from emb3d.compute import visualize
from emb3d.io import reader, writer
from emb3d.types import Backend, EmbedJob, ExecutionConfig

app = typer.Typer(add_completion=False)

# logging.basicConfig(level=logging.DEBUG)


def _input_file_or_stdin(input_file: Optional[Path], stdin_input: bool) -> TextIO:
    if stdin_input:
        # NOTE: This isn't memory efficient, stdin is primarily for convenience
        # Long jobs are better off using input_file flag
        return StringIO(sys.stdin.read())
    if input_file is None:
        input_file = Path(Prompt.ask("Enter the input file path"))
        if not input_file or not input_file.exists():
            raise typer.BadParameter(f"File {input_file} does not exist, aborting...")
    if not input_file.is_file():
        raise typer.BadParameter(f"File {input_file} not found, aborting...")
    return input_file.open()


def _output_file(
    out_file: Optional[Path], input_file: Optional[Path], stdin_input: bool
) -> TextIO:
    if out_file is not None:
        # TODO: Handle job termination/resume
        if out_file.exists():
            raise typer.BadParameter(f"File {out_file} already exists, aborting...")
        return out_file.open("w")
    elif stdin_input:
        return sys.stdout
    else:
        place_holder_suffix = input_file or Path("emb3d-run")
        default_out_file = place_holder_suffix.with_suffix(".out.jsonl")
        idx = 0
        while default_out_file.exists():
            typer.echo(
                f"Auto generating output file, file {default_out_file} already exists..."
            )
            idx += 1
            default_out_file = place_holder_suffix.with_suffix(f".{idx}.out.jsonl")
        return default_out_file.open("w")


def _pick_model(model_id: Optional[str]) -> str:
    default_model = config.AppConfig.instance().default_model
    return (
        model_id
        or default_model
        or Prompt.ask("Enter the embedding model", default="text-embedding-ada-002")
    )


def _execution_config(
    api_key: Optional[str], model_id: str, remote: bool
) -> ExecutionConfig:
    backend = EmbedJob.backend_from_model(model_id)
    remote_only_backends = (Backend.OPENAI, Backend.COHERE)
    if backend in remote_only_backends:
        remote = True

    if not remote:
        return ExecutionConfig.local()

    default_env_variables = {
        Backend.OPENAI: "OPENAI_API_KEY",
        Backend.COHERE: "CO_API_KEY",
        Backend.HUGGINGFACE: "HUGGINGFACE_API_KEY",
    }
    config_keys = {
        Backend.OPENAI: "openai_token",
        Backend.COHERE: "cohere_token",
        Backend.HUGGINGFACE: "huggingface_token",
    }
    cfg = config.AppConfig.instance()

    api_key = (
        api_key
        or cfg.get(config_keys[backend])
        or os.getenv(default_env_variables[backend])
    )

    if not api_key:
        raise typer.BadParameter(
            f"API key for {backend.value} backend is required, re-run the command with --api_key [your_key] or set {default_env_variables[backend]} environment variable."
        )

    return ExecutionConfig.remote(api_key=api_key)


@app.command("config", help="Get or set a configuration value.")
def cmd_config(
    key: str = typer.Argument(
        ...,
        help="The configuration key to set or get. ex: `default_model`, `openai_token` etc..",
    ),
    value: Optional[str] = typer.Argument(
        ...,
        help="The value associated with the key. If not provided, the current value will be printed.",
    ),
):
    cfg = config.AppConfig.instance()
    if value is None:
        current_value = cfg.get(key)
        if current_value is not None:
            typer.echo(f"{key} = {current_value}")
        else:
            typer.echo(f"{key} is not set.")
    else:
        cfg.set(key, value)
        typer.echo(f"Saved new value for {key} to emb3d config.")


def new_job_id(length: int = 5) -> str:
    """Generate a random string of the given length."""
    characters = string.hexdigits
    return "".join(random.choice(characters) for _ in range(length)).lower()


@app.command("compute", help="Compute embeddings for fun and profit.")
def cmd_compute(
    input_file: Optional[Path] = typer.Argument(
        ... if sys.stdin.isatty() else None,
        help="Path to the input file.",
    ),
    model: Optional[str] = typer.Option(
        config.AppConfig.instance().default_model,
        help="Embedding model to use.",
    ),
    output_file: Optional[Path] = typer.Option(
        None,
        "--output-file",
        "-o",
        help="Path to the output file. If not provided, a default path will be suggested.",
    ),
    api_key: Optional[str] = typer.Option(
        None,
        help="API key for the service hosting the model. If not provided, it will be prompted or fetched from environment variables.",
    ),
    remote: Annotated[
        bool,
        typer.Option(
            "--remote/--local",
            help="Choose whether to do inference locally or with an API token. This choice is available for sentence transformer and hugging face models. If a model cannot be run locally (ex: OpenAI models), this flag is ignored.",
        ),
    ] = True,
    batch_size: int = typer.Option(
        100,
        help="Batch size to use for inputs. Default is 100.",
    ),
    max_concurrent_requests: int = typer.Option(
        1000,
        help="(Remote Execution) Maximum number of concurrent requests for the embedding task. Default is 1000.",
    ),
):
    stdin_input = input_file is None

    input_file_io = _input_file_or_stdin(input_file, stdin_input)
    output_file_io = _output_file(output_file, input_file, stdin_input)
    model = _pick_model(model)
    execution_mode = _execution_config(api_key, model, remote)

    with input_file_io, output_file_io:
        num_records = sum(1 for _ in reader.line(input_file_io))
        # Rewind
        input_file_io.seek(0)
        new_job = EmbedJob(
            job_id=new_job_id(),
            in_file=input_file_io,
            model_id=model,
            out_file=output_file_io,
            total_records=num_records,
            batch_size=batch_size,
            max_concurrent_requests=min(max_concurrent_requests, num_records),
            execution_config=execution_mode,
        )

        compute.execute(new_job)


class ClusterOption(str, Enum):
    auto = "auto"
    cluster = "cluster"
    no_cluster = "no-cluster"


@app.command("visualize", help="Visualize generated embeddings.")
def cmd_visualize(
    embedding_file: Path = typer.Argument(
        ...,
        help="Path to the embedding file.",
    ),
    min_cluster_size: Optional[int] = typer.Option(
        None,
        help="Smallest size grouping that you wish to consider a cluster.",
    ),
    label_field: Optional[str] = typer.Option(
        "title",
        help="Field used to describe the record. If not provided, the `id` field (if present) or line numbers will be used.",
    ),
    cluster: Annotated[
        ClusterOption, typer.Option(case_sensitive=False)
    ] = ClusterOption.auto,
):
    with textui.SimpleProgressBar("Reading Data"):
        X, labels = visualize.get_data(embedding_file, label_field)

    n_records, n_dims = X.shape
    typer.echo(f"Loaded {n_records} records with {n_dims} dimensions.")
    with textui.SimpleProgressBar(
        f"Mapping records from {n_dims}-dimensional space to 2D (using: UMAP)."
    ):
        X_reduced = visualize.umap_reduce(X)

    needs_clustering = (
        cluster != ClusterOption.no_cluster
        and n_records > config.VISUALIZATION_CLUSTERING_THRESHOLD
    )

    if (
        n_records > config.VISUALIZATION_CLUSTERING_THRESHOLD
        and cluster != ClusterOption.cluster
    ):
        typer.echo(f"Too many records to visualize, clustering data...")

    min_cluster_size = min_cluster_size or config.VISUALIZATION_DEFAULT_MIN_CLUSTER_SIZE
    hdbscan_model = None
    if needs_clustering:
        with textui.SimpleProgressBar(
            f"Clustering with min_cluster_size {min_cluster_size} (using: HDSCAN)."
        ):
            hdbscan_model = visualize.cluster_hdbscan(X_reduced, min_cluster_size)

    with textui.SimpleProgressBar("Generating Scatter Plot"):
        chart = visualize.generate_chart(X_reduced, labels, hdbscan_model)
        out_file = embedding_file.with_suffix(".2d.html")
        writer.chart2html(chart, out_file)

    typer.echo(f"Visualization saved to {out_file}.")

    if typer.confirm("View in browser?"):
        webbrowser.open_new_tab("file://" + os.path.abspath(out_file))

emb3d/main.py
akhilravidas/emb3d
import io

from emb3d.types import EmbedJob, ExecutionConfig


def mock_embed_job(**kwargs):
    defaults = {
        "job_id": "test",
        "in_file": io.StringIO(),
        "model_id": "model",
        "out_file": io.StringIO(""),
        "total_records": 10,
        "batch_size": 10,
        "max_concurrent_requests": 1,
        "execution_config": ExecutionConfig.local(),
    }

    defaults.update(kwargs)
    return EmbedJob(**defaults)

emb3d/test_utils.py
akhilravidas/emb3d
"""
Rich UI widgets for job progress reporting
"""
import asyncio
import time

from rich.live import Live
from rich.panel import Panel
from rich.progress import BarColumn, Progress, ProgressColumn, SpinnerColumn, TextColumn
from rich.rule import Rule
from rich.table import Table, box
from rich.text import Text

from emb3d.types import EmbedJob, JobTracker

UI_UPDATE_INTERVAL = 0.8


class TaskProgressColumn(ProgressColumn):
    def __init__(self):
        super().__init__()

    def render(self, task) -> Text:
        max_width = len(str(task.total))
        completed_str = str(task.completed).rjust(max_width)
        total_str = str(task.total).rjust(max_width)
        return Text(f"{completed_str} / {total_str}")


class ProgressBar(Progress):
    def get_renderables(self):
        yield Rule("Records")
        overall_tasks = [
            task for task in self.tasks if task.fields.get("progress_type") == "overall"
        ]
        run_tasks = [
            task for task in self.tasks if task.fields.get("progress_type") != "overall"
        ]
        # Show overall progress at the end
        self.columns = (
            "{task.description}",
            BarColumn(),
            TaskProgressColumn(),
            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
        )
        yield self.make_tasks_table(run_tasks)
        self.columns = self.get_default_columns()
        yield Text("\n")
        yield Rule("Overall")
        yield self.make_tasks_table(overall_tasks)

    def update_values(self, tracker: JobTracker):
        for task in self.tasks:
            if task.fields.get("task_handle") == "processing":
                task.completed = tracker.encoding
            if task.fields.get("task_handle") == "failed":
                task.completed = tracker.failed
            if task.fields.get("task_handle") == "success":
                task.completed = tracker.saved
            if task.fields.get("task_handle") == "overall":
                # do the difference to get ETA
                self.advance(task.id, tracker.saved - task.completed)
        self.refresh()


def _recreate_progress(tracker: JobTracker):
    progress = ProgressBar()
    progress.add_task(
        "[magenta]Running",
        total=tracker.total,
        progress_type="other",
        task_handle="processing",
    )
    progress.add_task(
        "[red]Failed", total=tracker.total, progress_type="other", task_handle="failed"
    )
    progress.add_task(
        "[green]Saved",
        total=tracker.total,
        progress_type="other",
        task_handle="success",
    )
    progress.add_task(
        "Progress", total=tracker.total, progress_type="overall", task_handle="overall"
    )
    return progress


def render_ui_sync(job: EmbedJob, live: Live):
    progress = _recreate_progress(job.tracker)
    while job.tracker.saved < job.tracker.total:
        live.update(render_loop(job.tracker, progress))
        time.sleep(UI_UPDATE_INTERVAL)
    live.update(render_loop(job.tracker, progress))


async def render_ui_async(job: EmbedJob, live: Live):
    # TODO: Handle clean termination and keyboard interrupt
    progress = _recreate_progress(job.tracker)
    while job.tracker.saved < job.tracker.total:
        live.update(render_loop(job.tracker, progress))

        await asyncio.sleep(UI_UPDATE_INTERVAL)
    live.update(render_loop(job.tracker, progress))


def render_loop(tracker: JobTracker, progress: ProgressBar):
    table = Table.grid(expand=True)
    error_table = Table(show_header=False, box=box.SIMPLE, expand=True)
    for error in tracker.recent_errors:
        error_table.add_row("- " + str(error))

    progress.update_values(tracker)
    if tracker.recent_errors:
        table.add_row(
            Panel(
                error_table,
                title="[b] Recent Errors",
                title_align="left",
                border_style="red",
                padding=(1, 2),
                expand=True,
            )
        )
    table.add_row(
        Panel.fit(
            progress,
            title="[b]Jobs",
            title_align="left",
            border_style="green",
            padding=(1, 2),
        )
    )

    return table


class SimpleProgressBar:
    def __init__(self, text, transient=False):
        self.text = text
        self.progress = Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            transient=transient,
        )

    def __enter__(self):
        self.task_id = self.progress.add_task(self.text, total=None)
        self.progress.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.progress.stop()

emb3d/textui.py
akhilravidas/emb3d
"""
Type declarations
"""
from __future__ import annotations

from collections import deque
from dataclasses import dataclass, field
from enum import Enum
from typing import Dict, List, Optional, TextIO, Union


class Backend(Enum):
    """Supported model backends"""

    OPENAI = "OpenAI"
    COHERE = "Cohere"
    HUGGINGFACE = "Hugging Face"
    LOCAL = "Local Execution"


@dataclass
class Result:
    """Embedding call result wrapper for a batch"""

    data: List[List[float]]


@dataclass
class Failure:
    """Embedding call failure wrapper"""

    error: str


@dataclass
class WaitFor:
    """
    Embedding call rate limit response handler.
    Some services like hugging face provide a wait time before retrying.
    """

    seconds: float
    error: Optional[str]


EmbedResponse = Union[Result, Failure, WaitFor]

OpenAIModels = ("text-embedding-ada-002",)
CohereModels = (
    "embed-english-v2.0",
    "embed-english-light-v2.0",
    "embed-multilingual-v2.0",
)


@dataclass
class JobTracker:
    """
    Run stats for a single job
    """

    job_id: str
    success: int = 0
    encoding: int = 0
    failed: int = 0
    saved: int = 0
    total: int = 0
    recent_errors: deque = field(default_factory=lambda: deque(maxlen=5))


@dataclass
class ExecutionConfig:
    class ExecutionMode(Enum):
        REMOTE = "remote"
        LOCAL = "local"

    mode: ExecutionMode
    api_key: str

    @classmethod
    def local(cls) -> ExecutionConfig:
        return cls(cls.ExecutionMode.LOCAL, "")

    @classmethod
    def remote(cls, api_key: str) -> ExecutionConfig:
        return cls(cls.ExecutionMode.REMOTE, api_key)

    @property
    def is_remote(self) -> bool:
        return self.mode == self.ExecutionMode.REMOTE


@dataclass
class EmbedJob:
    """
    Single embedding job
    """

    job_id: str
    in_file: TextIO
    out_file: TextIO
    model_id: str
    total_records: int
    batch_size: int
    max_concurrent_requests: int
    execution_config: ExecutionConfig
    column_name: str = "text"
    tracker: JobTracker = field(init=False)

    def __post_init__(self):
        self.tracker = JobTracker(job_id=self.job_id, total=self.total_records)

    def batch_success(self, cnt: int):
        """Success callback"""
        self.tracker.success += cnt

    def batch_failure(self, cnt: int):
        """Failure callback"""
        self.tracker.failed += cnt

    def batch_saved(self, cnt: int):
        """Saved callback"""
        self.tracker.saved += cnt

    def batch_error(self, error_msg: str):
        """Error callback"""
        self.tracker.recent_errors.append(error_msg)

    @property
    def backend(self) -> Backend:
        """Model Backend"""
        return self.backend_from_model(self.model_id)

    @property
    def api_key(self) -> str:
        """API key"""
        return self.execution_config.api_key

    @classmethod
    def backend_from_model(cls, model_id: str) -> Backend:
        """model_id -> backend enum"""
        if model_id in OpenAIModels:
            return Backend.OPENAI
        elif model_id in CohereModels:
            return Backend.COHERE
        else:
            return Backend.HUGGINGFACE

    def describe(self) -> Dict[str, str]:
        return {
            "job_id": self.job_id,
            "model_id": self.model_id,
            "total_records": str(self.total_records),
            "batch_size": str(self.batch_size),
            "max_concurrent_requests": str(self.max_concurrent_requests),
            "mode": self.execution_config.mode.value,
        }


@dataclass
class Batch:
    """
    Input batch
    """

    row_ids: List[int]
    inputs: List[str]
    embeddings: Optional[List[List[float]]] = None
    error: Optional[str] = None

emb3d/types.py
akhilravidas/emb3d