data cleaning and curation for unstructured text
import sys
import os

sys.path.insert(0, os.path.abspath("../src"))

# Configuration file for the Sphinx documentation builder.
#
# For the full list of built-in configuration values, see the documentation:
# https://www.sphinx-doc.org/en/master/usage/configuration.html

# -- Project information -----------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information

project = "Galactic"
copyright = "Taylor AI, Inc., 2023"
author = "Benjamin Anderson & Brian Kim"

favicons = [
    "logo-dark.ico",
    "logo-light.png",
]

# -- General configuration ---------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration

extensions = [
    "sphinx.ext.autodoc",  # Automatic documentation from docstrings
    "sphinx.ext.viewcode",  # Add a link to the Python source code for classes, functions etc.
    "sphinx.ext.napoleon",  # Support for NumPy and Google style docstrings
    "sphinx_autodoc_typehints",  # Automatically document param types (less noise in class signature)
    "sphinx_copybutton",  # Add copy button to code blocks
    "sphinx_favicon",  # Add favicon
]

templates_path = ["_templates"]
exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"]
autodoc_typehints = "both"
autodoc_typehints_format = "fully-qualified"

highlight_language = "python"


html_theme_options = {
    "show_navbar_depth": 2,
    "repository_url": "https://github.com/taylorai/galactic",
    "use_issues_button": True,
    "icon_links": [
        {
            "name": "GitHub",
            "url": "https://github.com/taylorai/galactic",
            "icon": "fab fa-github-square",
        },
        {
            "name": "Slack",
            "url": "https://galactic-ai.slack.com/",
            "icon": "fab fa-slack",
        },
        {
            "name": "Twitter",
            "url": "https://twitter.com/TryTaylor_AI",
            "icon": "fab fa-twitter",
        },
        {
            "name": "LinkedIn",
            "url": "https://www.linkedin.com/company/taylor-ai/",
            "icon": "fab fa-linkedin",
        },
        {
            "name": "ProductHunt",
            "url": "https://www.producthunt.com/posts/taylor-ai",
            "icon": "fab fa-product-hunt",
        },
    ],
}

# -- Options for HTML output -------------------------------------------------
# https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output

html_theme = "sphinx_book_theme"
html_static_path = ["_static"]

docs/conf.py
taylorai/galactic
from .galactic import GalacticDataset

src/galactic/__init__.py
taylorai/galactic
### Code here adapted from openai cookbook: https://github.com/openai/openai-cookbook/blob/main/examples/api_request_parallel_processor.py
import aiohttp
import json
import numpy as np
import asyncio
import logging
import time
import tiktoken
from tqdm.auto import tqdm
from dataclasses import dataclass, field
from typing import Optional

logger = logging.getLogger("galactic")

EMBEDDING_URL = "https://api.openai.com/v1/embeddings"
CHAT_URL = "https://api.openai.com/v1/chat/completions"


@dataclass
class StatusTracker:
    num_tasks_started: int = 0
    num_tasks_in_progress: int = 0
    num_tasks_succeeded: int = 0
    num_tasks_failed: int = 0
    num_rate_limit_errors: int = 0
    time_of_last_rate_limit_error: int = 0
    total_requests = 0


@dataclass
class APIRequest:
    task_id: int
    type: str  # either embedding or chat
    text: str
    attempts_left: int
    system_prompt: Optional[str] = None
    logit_bias: Optional[dict[str, float]] = None
    max_new_tokens: Optional[int] = None
    result: list = field(default_factory=list)

    def __post_init__(self):
        # this is the same either way
        tokens = tiktoken.get_encoding("cl100k_base").encode(self.text)
        self.num_tokens = len(tokens)

        # get the URL and request JSON
        self.url = EMBEDDING_URL if self.type == "embedding" else CHAT_URL
        self.request_json = {}

        # you could save on API requests by doing > 1 text per request,
        # but we're usually hitting the token limit first anyway
        if self.type == "embedding":
            if self.logit_bias is not None or self.max_new_tokens is not None:
                raise NotImplementedError(
                    "Logit bias / max tokens doesn't apply to embeddings."
                )
            self.request_json["model"] = "text-embedding-ada-002"
            if len(tokens) > 8191:
                num_chunks = int(np.ceil(len(tokens) / 8191))
                self.request_json["input"] = np.array_split(
                    tokens, num_chunks
                ).tolist()
            else:
                self.request_json["input"] = [tokens]
        elif self.type == "chat":
            self.request_json["model"] = (
                "gpt-3.5-turbo"
                if self.num_tokens < 4000
                else "gpt-3.5-turbo-16k"
            )
            messages = []
            if self.system_prompt is not None:
                messages.append(
                    {"content": self.system_prompt, "role": "system"}
                )
            messages.append({"content": self.text, "role": "user"})
            self.request_json["messages"] = messages
            if self.logit_bias is not None:
                self.request_json["logit_bias"] = self.logit_bias
            if self.max_new_tokens is not None:
                self.request_json["max_tokens"] = self.max_new_tokens

    async def call_api(
        self,
        request_header: dict,
        retry_queue: asyncio.Queue,
        status_tracker: StatusTracker,
        pbar: Optional[tqdm] = None,
    ):
        try:
            status_tracker.total_requests += 1
            timeout = aiohttp.ClientTimeout(total=20)
            async with aiohttp.ClientSession(timeout=timeout) as session:
                async with session.post(
                    url=self.url,
                    headers=request_header,
                    json=self.request_json,
                ) as response:
                    response = await response.json()
            if "error" in response:
                logger.error(str(response))
                if "Rate limit" in response["error"].get("message", ""):
                    status_tracker.time_of_last_rate_limit_error = time.time()
                    status_tracker.num_rate_limit_errors += 1
                if "context length" in response["error"].get("message", ""):
                    logger.error(
                        "context length exceeded, retrying won't help"
                    )
                    self.attempts_left = 0
                self.result.append(response)
                if self.attempts_left:
                    self.attempts_left -= 1
                    retry_queue.put_nowait(self)
                else:
                    logger.error("out of tries")
                    status_tracker.num_tasks_in_progress -= 1
                    status_tracker.num_tasks_failed += 1
            else:
                if pbar is not None:
                    pbar.update(1)
                self.result.append(response)
                status_tracker.num_tasks_in_progress -= 1
                status_tracker.num_tasks_succeeded += 1
        except Exception as e:
            logger.error(f"Exception: {e}")
            self.result.append({"error": str(e)})
            if self.attempts_left:
                retry_queue.put_nowait(self)
            else:
                status_tracker.num_tasks_in_progress -= 1
                status_tracker.num_tasks_failed += 1


async def process_api_requests_from_list(
    texts: list[str],
    type: str,
    api_key: str,
    max_attempts: int,
    system_prompt: Optional[str] = None,
    logit_bias: Optional[dict[str, float]] = None,
    max_new_tokens: Optional[int] = None,
    max_tokens_per_minute: int = 90_000,
    max_requests_per_minute: int = 2000,
    show_progress: bool = True,
):
    if type not in ["embedding", "chat"]:
        raise ValueError("type must be either 'embedding' or 'chat'")

    """Processes API requests in parallel, throttling to stay under rate limits."""
    # constants
    seconds_to_pause_after_rate_limit_error = 15
    seconds_to_sleep_each_loop = (
        0.001  # 1 ms limits max throughput to 1,000 requests per second
    )

    # infer API endpoint and construct request header
    request_header = {"Authorization": f"Bearer {api_key}"}

    # initialize trackers
    queue_of_requests_to_retry = asyncio.Queue()
    status_tracker = (
        StatusTracker()
    )  # single instance to track a collection of variables
    next_request = None  # variable to hold the next request to call

    # initialize available capacity counts
    available_request_capacity = max_requests_per_minute
    available_token_capacity = max_tokens_per_minute
    last_update_time = time.time()

    # initialize flags
    texts_not_finished = True
    logger.debug(f"Initialization complete.")

    # turn the texts into an iterator
    if show_progress:
        pbar = tqdm(total=len(texts))
    else:
        pbar = None
    texts = iter(enumerate(texts))
    results = []
    while True:
        # get next request (if one is not already waiting for capacity)
        if next_request is None:
            if not queue_of_requests_to_retry.empty():
                next_request = queue_of_requests_to_retry.get_nowait()
                logger.debug(
                    f"Retrying request {next_request.task_id}: {next_request}"
                )
            elif texts_not_finished:
                try:
                    # get new request
                    idx, text = next(texts)
                    next_request = APIRequest(
                        task_id=idx,
                        type=type,
                        text=text,
                        attempts_left=max_attempts,
                        system_prompt=system_prompt,
                        logit_bias=logit_bias,
                        max_new_tokens=max_new_tokens,
                    )
                    status_tracker.num_tasks_started += 1
                    status_tracker.num_tasks_in_progress += 1
                    results.append(next_request)
                    logger.debug(
                        f"Reading request {next_request.task_id}: {next_request}"
                    )
                except StopIteration:
                    # if file runs out, set flag to stop reading it
                    logger.debug("Read file exhausted")
                    texts_not_finished = False

        # update available capacity
        current_time = time.time()
        seconds_since_update = current_time - last_update_time
        available_request_capacity = min(
            available_request_capacity
            + max_requests_per_minute * seconds_since_update / 60.0,
            max_requests_per_minute,
        )
        available_token_capacity = min(
            available_token_capacity
            + max_tokens_per_minute * seconds_since_update / 60.0,
            max_tokens_per_minute,
        )
        last_update_time = current_time

        # if enough capacity available, call API
        if next_request:
            next_request_tokens = next_request.num_tokens
            if (
                available_request_capacity >= 1
                and available_token_capacity >= next_request_tokens
            ):
                # update counters
                available_request_capacity -= 1
                available_token_capacity -= next_request_tokens
                next_request.attempts_left -= 1

                # call API
                asyncio.create_task(
                    next_request.call_api(
                        request_header=request_header,
                        retry_queue=queue_of_requests_to_retry,
                        status_tracker=status_tracker,
                        pbar=pbar,
                    )
                )
                logger.debug(
                    f"Called API for request {next_request.task_id}: {next_request}"
                )
                next_request = None  # reset next_request to empty

        # if all tasks are finished, break
        if status_tracker.num_tasks_in_progress == 0:
            break

        # main loop sleeps briefly so concurrent tasks can run
        await asyncio.sleep(seconds_to_sleep_each_loop)

        # if a rate limit error was hit recently, pause to cool down
        seconds_since_rate_limit_error = (
            time.time() - status_tracker.time_of_last_rate_limit_error
        )
        if (
            seconds_since_rate_limit_error
            < seconds_to_pause_after_rate_limit_error
        ):
            remaining_seconds_to_pause = (
                seconds_to_pause_after_rate_limit_error
                - seconds_since_rate_limit_error
            )
            await asyncio.sleep(remaining_seconds_to_pause)
            # ^e.g., if pause is 15 seconds and final limit was hit 5 seconds ago
            logger.warn(
                f"Pausing to cool down until {time.ctime(status_tracker.time_of_last_rate_limit_error + seconds_to_pause_after_rate_limit_error)}"
            )

    # after finishing, log final status
    logger.info(f"""Parallel processing complete.""")
    if status_tracker.num_tasks_failed > 0:
        logger.warning(
            f"{status_tracker.num_tasks_failed} / {status_tracker.num_tasks_started} requests failed."
        )
    if status_tracker.num_rate_limit_errors > 0:
        logger.warning(
            f"{status_tracker.num_rate_limit_errors} rate limit errors received. Consider running at a lower rate."
        )
    return results


def embed_texts_with_openai(
    texts: list[str],
    api_key: str,
    max_attempts: int = 3,
    max_tokens_per_minute: int = 350_000,
    max_requests_per_minute: int = 3500,
    show_progress: bool = True,
):
    results = asyncio.run(
        process_api_requests_from_list(
            texts=texts,
            type="embedding",
            api_key=api_key,
            max_attempts=max_attempts,
            max_requests_per_minute=max_requests_per_minute,
            max_tokens_per_minute=max_tokens_per_minute,
            show_progress=show_progress,
        )
    )
    # extract the embeddings
    embs = []
    for result in sorted(results, key=lambda x: x.task_id):
        if "error" in result.result[-1].keys():
            embs.append(None)
        else:
            data = result.result[-1]["data"]
            arr = np.array([r["embedding"] for r in data])
            avg = np.mean(arr, axis=0)
            embs.append((avg / np.linalg.norm(avg)).tolist())
    if len(embs) != len(texts):
        # write to json for debugging
        raise ValueError("Length of results does not match length of texts.")
    return embs


def run_chat_queries_with_openai(
    queries: list[str],
    api_key: str,
    system_prompt: Optional[str] = None,
    logit_bias: Optional[dict[str, float]] = None,
    max_new_tokens: Optional[int] = None,
    max_attempts: int = 3,
    max_tokens_per_minute: int = 90_000,
    max_requests_per_minute: int = 2000,
    show_progress: bool = True,
):
    results = asyncio.run(
        process_api_requests_from_list(
            texts=queries,
            type="chat",
            api_key=api_key,
            max_attempts=max_attempts,
            system_prompt=system_prompt,
            logit_bias=logit_bias,
            max_new_tokens=max_new_tokens,
            max_requests_per_minute=max_requests_per_minute,
            max_tokens_per_minute=max_tokens_per_minute,
            show_progress=show_progress,
        )
    )
    # extract the replies
    replies = [None for _ in range(len(queries))]
    for result in results:
        if isinstance(result.result[-1], str):
            logger.error(
                f"Result is a string instead of the expected dict: {result}"
            )
        if "error" in result.result[-1].keys():
            replies[result.task_id] = None
        else:
            replies[result.task_id] = result.result[-1]["choices"][0][
                "message"
            ]["content"]
    return replies

src/galactic/async_openai.py
taylorai/galactic
# pass in a subset of examples as few-shot to generate more, get back a new dataset with more examples

src/galactic/augment.py
taylorai/galactic
# This file exists purely to define a base class that helps a lot with IDE autocompletion/type hints. Don't worry about it.

from abc import ABC, abstractmethod, abstractclassmethod
from typing import (
    TypeVar,
    Optional,
    Union,
    Dict,
    List,
    Callable,
    Literal,
    Sequence,
    Any,
)
import pandas as pd
import numpy as np
import datasets
from .utils import (
    MessageConfig,
    OPENAI_MESSAGE_CONFIG,
)

T = TypeVar("T", bound="GalacticDatasetBase")


class GalacticDatasetBase(ABC):
    def __init__(self):
        pass

    ## loaders
    @abstractclassmethod
    def from_csv(cls: T, path: str) -> T:
        pass

    @abstractclassmethod
    def from_jsonl(cls: T, path: str) -> T:
        pass

    @abstractclassmethod
    def from_parquet(cls: T, path: str) -> T:
        pass

    @abstractclassmethod
    def from_pandas(cls: T, df: pd.DataFrame) -> T:
        pass

    @abstractclassmethod
    def from_hugging_face(
        cls: T,
        path: str,
        split: str,
        config_name: Optional[str] = None,
        **kwargs,
    ) -> T:
        pass

    @abstractclassmethod
    def from_hugging_face_stream(
        cls: T,
        path: str,
        split: str,
        config_name: Optional[str] = None,
        filters: list[Callable[[dict], bool]] = None,
        dedup_fields: Optional[list[str]] = None,
        max_samples: Optional[int] = 200000,
        **kwargs,
    ) -> T:
        pass

    @abstractmethod
    def save(self: T, path: str) -> None:
        pass

    ## conversation utils
    @abstractmethod
    def conversation_from_dicts(
        self,
        input_column: str,
        input_message_config: Union[MessageConfig, Dict],
        output_column: Optional[str] = None,
        output_message_config: MessageConfig = OPENAI_MESSAGE_CONFIG,
        conversation_field: Optional[str] = None,
        metadata_fields: Optional[List[str]] = None,
    ) -> T:
        pass

    @abstractmethod
    def conversation_from_string(
        self,
        input_column: str,
        user_delimiter: str,
        assistant_delimiter: str,
        system_delimiter: str,
        output_column: Optional[str] = None,
        output_message_config: MessageConfig = OPENAI_MESSAGE_CONFIG,
        strip_whitespace: bool = True,
        replace_values: Optional[Dict[str, str]] = None,
    ) -> T:
        pass

    @abstractmethod
    def convert_conversation_to_string(
        self,
        input_column: str,
        input_message_config: Union[MessageConfig, Dict],
        output_column: str,
        user_delimiter: str,
        assistant_delimiter: str,
        system_delimiter: str,
        strip_whitespace: bool = True,
    ) -> T:
        pass

    @abstractmethod
    def get_conversation_length(
        self, input_column: str, output_column: str
    ) -> T:
        pass

    @abstractmethod
    def get_conversation_speakers(
        self,
        input_column: str,
        input_message_config: Union[MessageConfig, Dict],
        output_column: str = "__speakers",
    ) -> T:
        pass

    @abstractmethod
    def is_alternating(
        self,
        input_column: str,
        input_message_config: Union[MessageConfig, Dict],
        output_column: str = "__alternating",
    ) -> T:
        pass

    @abstractmethod
    def get_last_speaker(
        self,
        input_column: str,
        input_message_config: Union[MessageConfig, Dict],
        output_column: str = "__last_speaker",
    ) -> T:
        pass

    @abstractmethod
    def standardize_last_turn(
        self,
        input_column: str,
        input_message_config: Union[MessageConfig, Dict],
        last_speaker_role: Literal["user", "assistant", "system"],
        output_column: Optional[str] = None,
    ) -> T:
        pass

    @abstractmethod
    def get_shared_prefix(
        self,
        input_column: str,
        input_message_config: Union[MessageConfig, Dict],
        output_column: str = "__shared_prefix",
    ) -> T:
        pass

    @abstractmethod
    def add_initial_system_message(
        self: T,
        input_column: str,
        input_message_config: Union[MessageConfig, dict],
        system_message: str,
        output_column: Optional[str] = None,
    ) -> T:
        pass

    @abstractmethod
    def take_initial_system_message(
        self,
        input_column: str,
        input_message_config: Union[MessageConfig, dict],
        output_column: str,
        output_type: Literal["dict", "string"] = "string",
        remove_from_input: bool = True,
    ) -> T:
        pass

    @abstractmethod
    def take_last_message(
        self: T,
        input_column: str,
        input_message_config: Union[MessageConfig, dict],
        output_column: str,
        output_type: Literal["dict", "string"] = "string",
        remove_from_input: bool = True,
    ) -> T:
        pass

    ## filters
    # @abstractmethod
    # def filter(
    #     self: T,
    #     fn: Callable,
    # ) -> T:
    #     pass

    # @abstractmethod
    # def filter_batches(
    #     self: T,
    #     fn: Callable[[List[Dict]], List[bool]]
    # ):
    #     pass

    @abstractmethod
    def apply_bloom_filter(
        self: T, fields: Sequence[str], inplace: bool = True
    ) -> Optional[T]:
        pass

    @abstractmethod
    def filter_string(
        self: T,
        fields: Sequence[str],
        values: Sequence[str],
        inplace: bool = True,
    ) -> Optional[T]:
        pass

    @abstractmethod
    def filter_regex(
        self: T, fields: Sequence[str], regex: str, inplace: bool = True
    ) -> Optional[T]:
        pass

    ## taggers
    # @abstractmethod
    # def tag(
    #     self: T,
    #     fn: Callable,
    #     fields: Sequence[str],
    #     inplace: bool = True
    # ) -> Optional[T]:
    #     pass

    @abstractmethod
    def tag_string(
        self, fields: Sequence[str], values: Sequence[str], tag: str
    ) -> T:
        pass

    @abstractmethod
    def tag_regex(self, fields: Sequence[str], regex: str, tag: str) -> T:
        pass

    @abstractmethod
    def detect_language(
        self,
        field: str,
        method: Literal["fasttext", "langdetect"] = "fasttext",
    ) -> T:
        pass

    @abstractmethod
    def detect_pii(
        self: T,
        fields: Sequence[str],
    ) -> T:
        pass

    @abstractmethod
    def count_tokens(
        fields: Sequence[str], tokenizer: Optional[str] = None
    ) -> T:
        pass

    @abstractmethod
    def calc_perplexity(
        self,
        field: str,
        model: str = "kenlm",  # other option is pythia
        language: Optional[str] = "en",
        dataset: Optional[str] = "wikipedia",
    ) -> T:
        pass

    @abstractmethod
    def ai_tagger(
        self: T,
        field: str,
        tags: Union[List[str], Dict[str, str]],
        prompt: Optional[str] = None,
        backend="openai",
        allow_not_sure: bool = False,
    ) -> T:
        pass

    ## transforms
    @abstractmethod
    def trim_whitespace(
        self: T, fields: Sequence[str], inplace: bool = True
    ) -> Optional[T]:
        pass

    @abstractmethod
    def unicode_normalize(
        self: T, fields: Sequence[str], inplace: bool = True
    ) -> Optional[T]:
        pass

    @abstractmethod
    def ai_column(
        self: T, fields: Sequence[str], inplace: bool = True
    ) -> Optional[T]:
        pass

    ## classifiers
    @abstractmethod
    def ai_classifier(
        self,
        new_column: str,
        field: Optional[str],
        classes: Union[list[str], dict[str, str]],
        prompt: Optional[str] = None,
        backend: Literal["openai", "huggingface", "embeddings"] = "openai",
    ) -> T:
        pass

    @abstractmethod
    def train_fasttext_classifier(
        self,
        model_name: str,
        save_dir: str,
        input_field: str,
        label_field: str,
        validation_split: float = 0.1,
        test_split: float = 0.1,
        normalize: list[str] = ["lower", "strip"],
        split_punctuation: bool = True,
        replace_newlines_with: str = " __newline__ ",
        target_model_size: str = "2M",
        training_duration: int = 300,
        random_seed: int = 42,
    ) -> None:
        pass

    @abstractmethod
    def train_embeddings_classifier(
        self,
        model_name: str,
        save_dir: str,
        label_field: str,
        model_type: str = "svm",
        input_field: str = "__embedding",
        validation_split=0.1,
        test_split=0.1,
        random_seed: int = 42,
    ) -> None:
        pass

    @abstractmethod
    def fasttext_classifier(
        self,
        new_column: str,
        model_path: str,
        field: str,
        # these should match how the model was trained
        normalize: list[str] = ["lower", "strip"],
        split_punctuation: bool = True,
        replace_newlines_with: str = " __newline__ ",
    ) -> T:
        pass

    @abstractmethod
    def embeddings_classifier(
        self, new_column: str, model_path: str, field: str = "__embedding"
    ) -> T:
        pass

    # ## embedding
    @abstractmethod
    def initialize_embedding_model(
        self: T, backend: Literal["auto", "cpu", "gpu", "openai"] = "auto"
    ) -> T:
        pass

    @abstractmethod
    def get_embeddings(
        self,
        input_field: str,
        embedding_field: str = "__embedding",
        backend: Literal["auto", "cpu", "gpu", "openai"] = "auto",
    ) -> T:
        pass

    @abstractmethod
    def get_nearest_neighbors(
        self,
        query: Union[str, np.ndarray, list[float]],
        k: int = 5,
        embedding_field: str = "__embedding",
    ) -> List[Dict[str, Any]]:
        pass

    @abstractmethod
    def reduce_embedding_dim(
        self,
        new_column: Optional[str] = None,
        method: Literal[
            "pca", "incremental_pca", "kernel_pca", "svd", "umap"
        ] = "pca",
        n_dims: int = 50,
        embedding_field: str = "__embedding",
        **kwargs,
    ) -> T:
        pass

    ## clustering
    @abstractmethod
    def cluster(
        self: T,
        n_clusters: Optional[int] = None,
        method: Literal[
            "kmeans", "minibatch_kmeans", "bisecting_kmeans", "hdbscan"
        ] = "kmeans",
        embedding_field: str = "__embedding",
        cluster_field: str = "__cluster",
        overwrite: bool = False,
        **kwargs,
    ):
        pass

    @abstractmethod
    def recompute_cluster_centers(
        self: T,
        cluster_field: str = "__cluster",
        embedding_field: str = "__embedding",
        method: Literal["centroid", "medoid"] = "centroid",
    ) -> T:
        pass

    @abstractmethod
    def _get_clusters(
        self,
        cluster_field: str = "__cluster",
    ) -> List[Dict[str, Any]]:
        pass

    @abstractmethod
    def remove_cluster(self: T, cluster_field: str, cluster_id: int) -> T:
        pass

    @abstractmethod
    def ai_label_clusters(
        self,
        new_column: str,
        context_fields: list[str],
        cluster_field: str = "__cluster",
        embedding_field: str = "__embedding",
        n_examples: int = 10,
        selection: Literal["random", "nearest"] = "random",
        prompt: Optional[str] = None,  # jinja2 template
    ) -> T:
        pass

    @abstractmethod
    def get_cluster_info(
        self,
        n_neighbors: int = 3,
        cluster_field: str = "__cluster",
        embedding_field: str = "__embedding",
        context_fields: list[str] = [],
        truncate_fields: int = 250,
        verbose: bool = True,
    ) -> T:
        pass

    @abstractmethod
    def semdedup(
        self,
        target_retention: Optional[float] = 0.8,
        threshold: Optional[float] = None,
        embedding_field: str = "__embedding",
        cluster_field: str = "__cluster",
        num_tuning_clusters: int = 3,
        dedup_strategy: Literal["random", "nearest", "furthest"] = "random",
        inplace=True,
    ):
        pass

    ## minhash lsh

    ## visualizations
    @abstractmethod
    def plot_embeddings(
        self,
        embedding_field: str,
        color_by: Optional[str] = None,
        save_path: Optional[str] = None,
        dot_size: int = 2,
        theme: Literal[
            "darkgrid", "whitegrid", "dark", "white", "ticks"
        ] = "whitegrid",
        width: Optional[int] = None,
        height: Optional[int] = None,
    ) -> None:
        pass

src/galactic/base.py
taylorai/galactic
import re
import os
import random
import tempfile
import fasttext
import numpy as np
import joblib
import tiktoken
import jinja2
from typing import Optional, Union, Literal
from sklearn.preprocessing import LabelEncoder
from sklearn.linear_model import SGDClassifier
from sklearn.model_selection import train_test_split
from .async_openai import run_chat_queries_with_openai

import logging

logger = logging.getLogger("galactic")


# Tendency is that bigger target sizes take longer to optimize, so if you only have 5 min use 2M or 1M or 500K
def train_fasttext_classifier(
    self,
    model_name: str,
    save_dir: str,
    input_field: str,
    label_field: str,
    validation_split: float = 0.1,
    test_split: float = 0.1,
    normalize: list[str] = ["lower", "strip"],
    split_punctuation: bool = True,
    replace_newlines_with: str = " __newline__ ",
    target_model_size: str = "2M",
    training_duration: int = 300,
    random_seed: int = 42,
):
    """
    Trains a fasttext classifier on the dataset provided to the instance.

    Args:
        self: The instance of the class.
        model_name (str): The name of the model to be saved.
        save_dir (str): The directory where the model should be saved.
        input_field (str): The field name from the dataset to be used as input.
        label_field (str): The field name from the dataset to be used as label.
        validation_split (float, optional): The proportion of the dataset to include in the validation split. Defaults to 0.1.
        test_split (float, optional): The proportion of the dataset to include in the test split. Defaults to 0.1.
        normalize (list[str], optional): List of normalization to apply to the text. Defaults to ["lower", "strip"].
        split_punctuation (bool, optional): Whether to split the text by punctuation. Defaults to True.
        replace_newlines_with (str, optional): String to replace newlines with in the text. Defaults to " __newline__ ".
        target_model_size (str, optional): The desired size of the fasttext model. Defaults to "2M".
        training_duration (int, optional): The duration in seconds for the fasttext autotune. Defaults to 300.
        random_seed (int, optional): Seed for reproducing results. Defaults to 42.
    """

    def _preprocess(text):
        text = str(text)
        for fn in normalize:
            text = getattr(str, fn)(text)
        text = text.replace("\n", replace_newlines_with)
        for n in normalize:
            if n == "lower":
                text = text.lower()
            elif n == "strip":
                text = text.strip()
        if split_punctuation:
            text = re.sub(r"([.!?,\'/()])", r" \1 ", text)
        return text

    def _get_label(sample):
        return "__label__" + str(sample[label_field]).lower().replace(" ", "_")

    def _get_line(sample):
        return _get_label(sample) + " " + _preprocess(sample[input_field])

    def _create_train_files(dataset):
        random.seed(random_seed)
        train_file = tempfile.NamedTemporaryFile(mode="a", delete=False)
        val_file = tempfile.NamedTemporaryFile(mode="a", delete=False)
        test_file = tempfile.NamedTemporaryFile(mode="a", delete=False)
        for sample in dataset:
            if random.random() < validation_split:
                val_file.write(_get_line(sample) + "\n")
            elif random.random() < test_split + validation_split:
                test_file.write(_get_line(sample) + "\n")
            else:
                train_file.write(_get_line(sample) + "\n")
        # return names of files used for training
        return train_file.name, val_file.name, test_file.name

    # create train and validation files
    train_file, val_file, test_file = _create_train_files(self.dataset)

    # train fasttext model
    model = fasttext.train_supervised(
        input=train_file,
        autotuneValidationFile=val_file,
        autotuneDuration=training_duration,
        autotuneModelSize=target_model_size,
    )

    # test model on test set, report individual class scores
    overall = model.test(test_file)
    scores = model.test_label(test_file)
    scores = {k: v["precision"] for k, v in scores.items()}
    logger.info(f"Test set accuracy: {overall[1]}")
    logger.info(f"Test set accuracy per-class: {scores}")

    # save model
    model.save_model(os.path.join(save_dir, model_name + ".ftz"))


def train_embeddings_classifier(
    self,
    model_name: str,
    save_dir: str,
    label_field: str,
    model_type: str = "svm",
    input_field: str = "__embedding",
    validation_split=0.1,
    test_split=0.1,
    random_seed: int = 42,
):
    """
    Trains a classifier on the dataset using embeddings.

    Args:
        self: The instance of the class.
        model_name (str): The name of the model to be saved.
        save_dir (str): The directory where the model should be saved.
        label_field (str): The field name from the dataset to be used as label.
        model_type (str, optional): The type of model to train. Either "svm" or "logistic". Defaults to "svm".
        input_field (str, optional): The field name from the dataset to be used as input. Defaults to "__embedding".
        validation_split (float, optional): The proportion of the dataset to include in the validation split. Defaults to 0.1.
        test_split (float, optional): The proportion of the dataset to include in the test split. Defaults to 0.1.
        random_seed (int, optional): Seed for reproducing results. Defaults to 42.
    """
    if "__embedding" not in input_field:
        logger.warning(
            "This is designed to use the embeddings (typically '__embedding') as input. I hope you know what you're doing..."
        )
    elif input_field not in self.dataset.features:
        raise ValueError(f"Input field {input_field} not found in dataset.")
    elif label_field not in self.dataset.features:
        raise ValueError(f"Label field {label_field} not found in dataset.")
    # TODO: make sure input_field is a list of numbers to avoid footgun
    X = np.array(self.dataset[input_field])

    # encode labels
    le = LabelEncoder()
    y = le.fit_transform(self.dataset[label_field])

    # split train & test
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=test_split, random_state=random_seed
    )

    # train model
    if model_type not in ["svm", "logistic"]:
        raise ValueError(
            f"Model type {model_type} not supported. Use 'svm' or 'logistic'."
        )
    loss = "hinge" if model_type == "svm" else "log_loss"
    model = SGDClassifier(
        random_state=random_seed,
        loss=loss,
        validation_fraction=validation_split,
        early_stopping=True,
    )
    model.fit(X_train, y_train)

    # display train & test accuracy
    train_acc = model.score(X_train, y_train)
    test_acc = model.score(X_test, y_test)
    logger.info(f"Training accuracy: {train_acc}")
    logger.info(f"Test accuracy: {test_acc}")
    logger.info(
        f"Validation split was used for early stopping, not for measuring out-of-sample accuracy."
    )

    # save model & label encoder
    joblib.dump(model, os.path.join(save_dir, model_name + ".joblib"))
    joblib.dump(le, os.path.join(save_dir, model_name + ".labels.joblib"))


def ai_classifier(
    self,
    new_column: str,
    field: Optional[str],
    classes: Union[list[str], dict[str, str]],
    prompt: Optional[str] = None,
    backend: Literal["openai", "huggingface", "embeddings"] = "openai",
):
    """
    Classify a field using OpenAI's API or a HF zero-shot classifier to generate a new column based on an existing column.

    Args:
        self: The instance of the class.
        new_column (str): Name of the new column to be added to the dataset.
        field (Optional[str]): Name of the field to use for classification. Can be None if backend is 'embeddings'.
        classes (Union[list[str], dict[str, str]]): Classes to classify into. Can be a list of labels or a dict of label: description.
        prompt (Optional[str], optional): Prompt template (Jinja2) to use for the API request, with a template slot for the field to classify. Defaults to None.
        backend (str, optional): Which backend to use. Currently supports 'embeddings', 'openai', and 'huggingface'. Defaults to "openai".
    """
    if backend == "embeddings":
        # warn if field is specified that embeddings ignores the field
        if field is not None:
            logger.warning(
                f"Field {field} specified, but we're using embeddings to classify, which will use whatever field was embedded with get_embeddings()."
            )
        # make sure that we've already embedded the dataset
        if "__embedding" not in self.dataset.column_names:
            raise RuntimeError(
                "Dataset does not have embeddings. Run get_embeddings() first."
            )
        elif self.model is None:
            raise RuntimeError(
                "Dataset does not have an embedding model. Run get_embeddings() first, or run get_embedding_model() with the backend that matches '__embeddings'."
            )
        # embed the classes
        idx2label = (
            {idx: label for idx, label in enumerate(classes)}
            if isinstance(classes, list)
            else {idx: label for idx, label in enumerate(classes.keys())}
        )
        idx2input = (
            idx2label
            if isinstance(classes, list)
            else {idx: classes[label] for idx, label in idx2label.items()}
        )
        class_embeddings = np.array(
            [self.model(idx2input[i]) for i in range(len(idx2input))]
        )

        # get the embeddings for the dataset
        self.dataset = self.dataset.map(
            lambda sample: {
                new_column: idx2label[
                    np.dot(class_embeddings, sample["__embedding"]).argmax()
                ]
            }
        )

    elif backend == "openai":
        if self.openai_api_key is None:
            self.openai_api_key = os.environ.get("OPENAI_API_KEY", None)
            if self.openai_api_key is None:
                raise ValueError(
                    "No OpenAI API key found. Set OPENAI_API_KEY environment variable, or set the openai_api_key attribute of the GalacticDataset."
                )
        # make sure field exists
        if field not in self.dataset.column_names:
            raise ValueError(f"Field {field} not found in dataset.")

        # if prompt is not provided, make default prompt based on class descriptions
        if prompt is None:
            prompt = "Classify the provided text into one of the following classes:\n\n"
            if isinstance(classes, list):
                for label in classes:
                    prompt += f"- {label}\n"
            else:
                for label, description in classes.items():
                    prompt += f"- {label}: {description}\n"
            prompt += "\n---\n\nText: {{" + field + "}}\n\n---\n\nClass:"

        # create logit bias
        token2cls = {}
        tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")
        class_list = (
            list(classes)
            if isinstance(classes, list)
            else list(classes.keys())
        )
        for cls in class_list:
            t1 = tokenizer.encode(cls)[0]
            t2 = tokenizer.encode(" " + cls)[0]
            if t1 in token2cls or t2 in token2cls:
                raise ValueError(
                    "Class names must not share a prefix for logit bias trick to work. Try changing the labels and see if that helps."
                )
            else:
                token2cls[t1] = cls
                token2cls[t2] = cls
        logit_bias = {t: 100 for t in token2cls.keys()}

        # create queries
        template = jinja2.Template(prompt)
        prompts = self.dataset.select_columns([field]).map(
            lambda sample: {"__prompt": template.render(**sample)}
        )["__prompt"]
        logger.info(f"Example prompt: {prompts[0]}")

        # run queries
        responses = run_chat_queries_with_openai(
            queries=prompts,
            api_key=self.openai_api_key,
            logit_bias=logit_bias,
            max_new_tokens=1,
            max_tokens_per_minute=self.max_tokens_per_minute,
            max_requests_per_minute=self.max_requests_per_minute,
        )

        # map back to labels
        first_tokens = [tokenizer.encode(response) for response in responses]
        labels = [token2cls[t[0]] for t in first_tokens]
        self.dataset = self.dataset.add_column(name=new_column, column=labels)

    elif backend == "huggingface":
        from transformers import pipeline

        pipe = pipeline(
            "zero-shot-classification",
            model="mrm8488/deberta-v3-small-finetuned-mnli",
        )
        if isinstance(classes, list):
            candidate_answers = classes
            ans2label = {c: c for c in classes}
        else:
            candidate_answers = list(classes.values())
            ans2label = {classes[k]: k for k in classes.keys()}

        self.dataset = self.dataset.map(
            lambda sample: {
                new_column: ans2label[
                    pipe(sample[field], candidate_answers)["labels"][0]
                ]
            }
        )
    else:
        raise ValueError(f"Unknown backend: {backend}")

    return self


def fasttext_classifier(
    self,
    new_column: str,
    model_path: str,
    field: str,
    # these should match how the model was trained
    normalize: list[str] = ["lower", "strip"],
    split_punctuation: bool = True,
    replace_newlines_with: str = " __newline__ ",
):
    """
    Classify a field using a trained fastText model to generate a new column based on an existing column.

    Args:
        self: The instance of the class.
        fasttext_model (str): The path to the trained fastText model file.
        new_column (str): Name of the new column to be added to the dataset.
        field (Optional[str]): Name of the field to use for classification. Can be None if the fastText model is trained on embeddings.
        k (int, optional): Number of labels to predict. Defaults to 1.
    """
    # make sure the input field is a string
    if self.dataset.features[field].dtype != "string":
        raise ValueError(f"Field {field} is not a string field.")

    model = fasttext.load_model(model_path)

    def _preprocess(text):
        for fn in normalize:
            text = getattr(str, fn)(text)
        text = text.replace("\n", replace_newlines_with)
        for n in normalize:
            if n == "lower":
                text = text.lower()
            elif n == "strip":
                text = text.strip()
        if split_punctuation:
            text = re.sub(r"([.!?,\'/()])", r" \1 ", text)
        return text

    def _classify(sample):
        input_text = _preprocess(sample[field])
        output = model.predict(input_text)
        return {new_column: output[0][0].split("__label__")[1]}

    self.dataset = self.dataset.map(_classify)
    logger.info(
        f"Applied classifier to field {field}, added result to {new_column}."
    )
    return self


def embeddings_classifier(
    self,
    new_column: str,
    model_path: str,
    field: str = "__embedding",
):
    """Use an embeddings model to classify a field."""
    model = joblib.load(model_path + ".joblib")
    le = joblib.load(model_path + ".labels.joblib")

    def _classify(sample):
        outputs = le.inverse_transform(model.predict(sample[field]))
        return {new_column: outputs}

    self.dataset = self.dataset.map(_classify, batched=True)
    logger.info(
        f"Applied classifier to field {field}, added result to {new_column}."
    )

    return self

src/galactic/classifiers.py
taylorai/galactic
import datasets
import networkx as nx
import numpy as np
import random
from typing import Optional, Literal
import numpy as np
from sklearn.cluster import MiniBatchKMeans, KMeans, BisectingKMeans, HDBSCAN
from collections import Counter
import jinja2
from .async_openai import run_chat_queries_with_openai

import logging

logger = logging.getLogger("galactic")


def cluster(
    self,
    n_clusters: Optional[int] = None,
    method: Literal[
        "kmeans", "minibatch_kmeans", "bisecting_kmeans", "hdbscan"
    ] = "kmeans",
    embedding_field: str = "__embedding",
    cluster_field: str = "__cluster",
    overwrite: bool = False,
    **kwargs,
):
    """
    Cluster the dataset using the specified method.

    .. code-block:: python

        # Cluster your dataset into 10 clusters using minibatch k-means
        ds.cluster(n_clusters=10)

    :param n_clusters: Required. The number of clusters to form.
    :param method: Optional. The clustering method to use. Default = ``kmeans``.
    :param embedding_field: Optional. Specify a name for the field to use for clustering. Default = ``__embedding``.
    :param kwargs: Optional. Additional keyword arguments to pass to the clustering algorithm.

    """

    if embedding_field not in self.dataset.column_names:
        raise ValueError(
            "You must call get_embeddings() before calling cluster(). If your dataset already has an embeddings column, pass it as 'embedding_field' argument."
        )
    # check if cluster_field is already set
    if cluster_field in self.dataset.column_names:
        if overwrite:
            logger.warning(
                f"You already have clusters in field {cluster_field}. Since overwrite=True, these will be overwritten."
            )
            self.dataset = self.dataset.select_columns(
                [c for c in self.dataset.column_names if c != cluster_field]
            )
        else:
            raise ValueError(
                f"You already have clusters in field {cluster_field}. If you want to overwrite them, pass overwrite=True. Otherwise, use a different 'cluster_field' to create a new clustering."
            )

    # set cluster ids and centers, in a way that allows multiple clusterings!
    if self.cluster_ids is None:
        self.cluster_ids = {}
    if self.cluster_centers is None:
        self.cluster_centers = {}

    # check if emb dimension is large
    if len(self.dataset[embedding_field][0]) >= 384:
        logger.info(
            "Embedding dimension is large, which is fine! But consider also experimenting with dimensionality reduction before clustering."
        )
    if method != "hdbscan" and n_clusters is None:
        raise ValueError(
            "You must specify the number of clusters with n_clusters argument. If you don't want to set this a priori, try using the hdbscan method, which doesn't require you to set it."
        )

    if method == "minibatch_kmeans":
        model = MiniBatchKMeans(
            n_clusters=n_clusters, batch_size=kwargs.get("batch_size", 1024)
        )
        for epoch in range(kwargs.get("n_epochs", 5)):
            logger.info(f"Epoch {epoch+1}/{kwargs.get('n_epochs', 5)}")
            self.dataset.map(
                lambda x: model.partial_fit(np.array(x[embedding_field])),
            )
        self.dataset = self.dataset.map(
            lambda batch: {
                cluster_field: model.predict(
                    batch[embedding_field],
                )
            },
            batched=True,
            batch_size=kwargs.get("batch_size", 1024),
        )

    elif method in ["kmeans", "bisecting_kmeans"]:
        if method == "kmeans":
            model = KMeans(n_clusters=n_clusters, init="k-means++", n_init=1)
        elif method == "bisecting_kmeans":
            model = BisectingKMeans(
                n_clusters=n_clusters,
                init="k-means++",
                n_init=1,
                bisecting_strategy="largest_cluster",
            )
        arr = np.array(self.dataset[embedding_field])
        labels = model.fit_predict(arr)
        self.cluster_ids[cluster_field] = list(set(labels))
        # cluster centers is a dict of id -> center
        self.cluster_centers[cluster_field] = {
            i: model.cluster_centers_[i]
            for i in range(max(self.cluster_ids[cluster_field]))
        }
        self.dataset = self.dataset.add_column(cluster_field, labels)

    elif method == "hdbscan":
        min_cluster_size = kwargs.get("min_cluster_size", 25)
        model = HDBSCAN(
            min_cluster_size=min_cluster_size,
            min_samples=kwargs.get("min_samples", min_cluster_size),
            alpha=kwargs.get("alpha", 1.0),
            cluster_selection_method=kwargs.get(
                "cluster_selection_method", "eom"
            ),
            leaf_size=kwargs.get("leaf_size", 40),
            store_centers="medoid",
        )
        arr = np.array(self.dataset[embedding_field])
        labels = model.fit_predict(arr)
        self.cluster_ids[cluster_field] = list(set(labels))
        # cluster centers is a dict of id -> center
        self.cluster_centers[cluster_field] = {
            i: model.medoids_[i] for i in range(max(self.cluster_ids))
        }
        self.dataset = self.dataset.add_column(cluster_field, labels)
    else:
        raise ValueError(f"Unknown clustering method: {method}")


def recompute_cluster_centers(
    self,
    cluster_field: str = "__cluster",
    embedding_field: str = "__embedding",
    method: Literal["centroid", "medoid"] = "centroid",
):
    """
    Re-compute centers and store them in self.cluster_centers. There are two main reasons to do this:
    1. You've deduplicated or sparsified clusters, so their centers have moved, or
    2. You've re-loaded the data from file, so the centers are no longer in memory.
    """
    pass


# preferred to filtering out the cluster, because it will remove the cluster from the cluster_ids list
def remove_cluster(self, cluster_field: str, cluster_id: int):
    """Remove a cluster from the dataset."""
    if cluster_field not in self.dataset.column_names:
        raise ValueError(
            f"Cluster field {cluster_field} not found in dataset."
        )
    else:
        self.dataset = self.dataset.filter(
            lambda x: x[cluster_field] != cluster_id
        )
        self.cluster_ids[cluster_field].remove(cluster_id)
        del self.cluster_centers[cluster_field][cluster_id]
    logger.info(
        f"Removed cluster {cluster_id} from clustering {cluster_field} from dataset."
    )


# this does loop through the whole dataset, but that's faster than looping through it N times for N clusters
def _get_clusters(self, cluster_field: str):
    clusters = {}
    self.dataset.map(
        lambda x: clusters.setdefault(x[cluster_field], []).append(x),
    )
    return {
        cluster_id: datasets.Dataset.from_list(cluster)
        for cluster_id, cluster in clusters.items()
    }


def _get_nearest_in_cluster(
    cluster: datasets.Dataset,
    embedding_field: str,
    query: np.ndarray,
    k: int = 3,
):
    """Get the nearest neighbors of a query point in a cluster."""
    emb_matrix = np.array(cluster[embedding_field])
    similarities = np.dot(emb_matrix, query)
    top_k = list(np.argsort(similarities)[::-1][:k])
    examples = list(cluster.select(top_k))
    return examples


def ai_label_clusters(
    self,
    new_column: str,
    context_fields: list[str],
    cluster_field: str = "__cluster",
    embedding_field: str = "__embedding",
    n_examples: int = 10,
    selection: Literal["random", "nearest"] = "random",
    prompt: Optional[str] = None,  # jinja2 template
):
    """
    Labels the clusters using AI based on the given fields and prompt template.

    .. code-block:: python

        # Example usage:
        ds.ai_label_clusters(fields=['field1', 'field2'], selection='nearest')

    :param fields: List of fields used to identify clusters.
    :param new_column: Optional. Name of the new column where the cluster labels will be stored. Defaults to '__cluster_label'.
    :param n_examples: Optional. Number of examples to consider for labeling. Defaults to 10.
    :param selection: Optional. Strategy to select examples for labeling. Can be 'random' or 'nearest'. Defaults to 'random'.
    :param embedding_field: Optional. Name of the embedding field to be used. Defaults to '__embedding'.
    :param prompt: Optional. Jinja2 template string to be used as the prompt for labeling.
    :return: The modified object with labeled clusters.
    """

    if not prompt:
        # Default Jinja2 template
        prompt = (
            "Please identify a single shared topic or theme among the following examples in a few words. "
            "It's ok if there are a small minority of examples that don't fit with the theme, but if "
            "there's no clear shared topic or theme, just say 'No shared topic or theme'.\n\n"
            "{% for example in examples %}\n"
            "\t### Example {{ loop.index }}\n"
            "\t{% for field in fields %}\n"
            "\t\t- {{ field }}: {{ example[field] }}\n"
            "\t{% endfor %}\n"
            "{% endfor %}\n\n"
            "Now, state the single topic or theme in 3-10 words, no long lists:"
        )
    template = jinja2.Template(prompt)
    queries = []

    logger.info("Splitting dataset into clusters... (this might take a bit).")
    clusters = self._get_clusters(cluster_field)
    ids = []
    for id, cluster in clusters.items():
        ids.append(id)
        cluster_center = self.cluster_centers[cluster_field][id]
        cluster = cluster.select_columns(context_fields + [embedding_field])
        if len(cluster) < n_examples:
            examples = list(cluster.select_columns(context_fields))
        elif selection == "nearest":
            examples = _get_nearest_in_cluster(
                cluster, embedding_field, cluster_center, k=n_examples
            ).select_columns(context_fields)
        elif selection == "random":
            examples = random.choices(
                list(cluster.select_columns(context_fields)), k=n_examples
            )
        else:
            raise ValueError(f"Unknown selection method: {selection}")

        prompt = template.render(examples=examples, fields=context_fields)
        queries.append(prompt)

    responses = run_chat_queries_with_openai(queries, self.openai_api_key)
    id2label = {id: response for id, response in zip(ids, responses)}
    self.dataset = self.dataset.map(
        lambda x: {new_column: id2label[x[cluster_field]]}
    )
    return self


def get_cluster_info(
    self,
    n_neighbors: int = 3,
    cluster_field: str = "__cluster",
    embedding_field: str = "__embedding",
    context_fields: list[str] = [],
    truncate_fields: int = 250,
    verbose: bool = True,
):
    """
    Retrieves information regarding clusters and their nearest neighbors.

    .. code-block:: python

        # Example usage:
        ds.get_cluster_info(n_neighbors=5, field='title')

    :param n_neighbors: Number of nearest neighbors to be retrieved for each cluster. Defaults to 3.
    :param field: Optional specific field to be printed from the neighbors.
    :raises ValueError: If cluster() or get_embeddings() methods are not called prior to calling this method.
    """
    result = []
    if (
        not hasattr(self, "cluster_centers")
        or len(self.cluster_centers.keys()) == 0
    ):
        raise ValueError(
            "You must call cluster() before calling get_cluster_info()"
        )
    if not hasattr(self, "model"):
        raise ValueError(
            "You must call get_embeddings() before calling get_cluster_info(). If your dataset already has embeddings, you need to re-initialize the embedding model with initialize_embedding_model() so that new embeddings can be computed for queries."
        )
    if cluster_field not in self.dataset.column_names:
        raise ValueError(
            f"Cluster field {cluster_field} not found in dataset."
        )

    counts = Counter(self.dataset[cluster_field])

    # for each one, get the 3 nearest neighbors
    clusters = self._get_clusters(cluster_field)
    for id, cluster in clusters.items():
        cluster_center = self.cluster_centers[cluster_field][id]
        examples = _get_nearest_in_cluster(
            cluster, embedding_field, cluster_center, k=n_neighbors
        )

        cluster_info = {
            "cluster_id": id,
            "cluster_size": counts[id],
            # keep only fields in context_fields
            "examples": [
                {field: example[field] for field in context_fields}
                for example in examples
            ],
        }
        result.append(cluster_info)
    if verbose:
        for cluster_info in result:
            print(
                f"Cluster {cluster_info['cluster_id']} ({cluster_info['cluster_size']} items)"
            )
            if len(context_fields) > 0:
                for idx, example in enumerate(cluster_info["examples"]):
                    print(f"### Example {idx + 1}")
                    for field in context_fields:
                        print(
                            f"\t- {field}: {example[field][:truncate_fields]}"
                        )
    return result


def get_duplicates(
    cluster: datasets.Dataset,
    cluster_center: np.ndarray,
    threshold: float,
    emb_matrix: Optional[
        np.ndarray
    ] = None,  # pass this in if you already have it
    similarities: Optional[
        np.ndarray
    ] = None,  # pass this in if you already have it
    embedding_field: str = "__embedding",
    dedup_strategy: Literal["random", "nearest", "furthest"] = "random",
):
    """
    Retrieves duplicates within a specified cluster based on a threshold.

    .. code-block:: python

        # Example usage:
        duplicates = get_duplicates(cluster=ds, threshold=0.8, strategy='nearest')

    :param cluster: The dataset cluster to be analyzed for duplicates.
    :param threshold: Similarity threshold to consider two items as duplicates.
    :param strategy: Strategy to select duplicates, can be 'random', 'nearest', or 'furthest'. Defaults to 'random'.
    :raises ValueError: If an unknown strategy is provided.
    :return: List of identified duplicate items within the cluster.
    """

    duplicates = []
    num_points = len(cluster)
    if emb_matrix is None:
        emb_matrix = np.array(cluster[embedding_field])
    if dedup_strategy != "random":
        # centroid = np.mean(emb_matrix, axis=0)
        # no longer doing this because for e.g. hdbscan we might use the medoid instead.
        # center should be passed as an argument
        id2dist = {
            cluster[i]["__id"]: np.dot(emb_matrix[i], cluster_center)
            for i in range(num_points)
        }

    # find connected components
    if similarities is None:
        similarities = np.dot(emb_matrix, emb_matrix.T)
    G = nx.Graph()
    for i in range(num_points):
        for j in range(i + 1, num_points):
            if similarities[i][j] > threshold:
                G.add_edge(
                    cluster[i]["__id"],
                    cluster[j]["__id"],
                )
    # get duplicates
    for cmp in nx.connected_components(G):
        cmp = list(cmp)
        if dedup_strategy == "nearest":
            cmp.sort(key=lambda x: id2dist[x])
        elif dedup_strategy == "furthest":
            cmp.sort(key=lambda x: -id2dist[x])
        duplicates.extend(cmp[1:])
    return duplicates


def tune_threshold(
    cluster: datasets.Dataset,
    cluster_center: np.ndarray,
    target_retention: float,
    tol: float = 0.01,
    max_iter: int = 30,
    embedding_field: str = "__embedding",
    dedup_strategy: Literal["random", "nearest", "furthest"] = "random",
):
    """
    Tunes the threshold for identifying duplicates within a cluster to achieve a target retention rate.

    .. code-block:: python

        # Example usage:
        threshold = tune_threshold(cluster=ds, target_retention=0.9)

    :param cluster: The dataset cluster to be analyzed.
    :param target_retention: Target retention rate to achieve.
    :param tol: Tolerance for the difference between achieved and target retention rate. Defaults to 0.01.
    :param max_iter: Maximum number of iterations for tuning. Defaults to 30.
    :return: The tuned threshold value.
    """

    tol = max(tol, 1 / len(cluster))
    emb_matrix = np.array(cluster[embedding_field])
    similarities = np.dot(emb_matrix, emb_matrix.T)
    min_sim = np.min(similarities)
    max_sim = np.max(similarities)

    # use binary search to find threshold
    print("Minimum similarity:", min_sim, "Maximum similarity:", max_sim)
    lo = min_sim
    hi = max_sim

    for _ in range(max_iter):
        mid = (lo + hi) / 2
        duplicates = get_duplicates(
            cluster=cluster,
            cluster_center=cluster_center,
            threshold=mid,
            emb_matrix=emb_matrix,
            similarities=similarities,
            embedding_field=embedding_field,
            dedup_strategy=dedup_strategy,
        )
        # Calculate the retention rate after removing duplicates
        retention = 1 - len(duplicates) / len(cluster)
        print(f"Threshold: {round(mid, 3)}, Retention: {round(retention, 3)}")

        # if retention is within tolerance, we're done
        if abs(retention - target_retention) < tol:
            return mid
        # if retention is too low, increase the threshold (less filtering)
        elif retention < target_retention:
            lo = mid
        # if retention is too high, lower the threshold (more filtering)
        elif retention > target_retention:
            hi = mid

    # Final threshold is the average of the final bounds
    final_threshold = (hi + lo) / 2
    return final_threshold


def semdedup(
    self,
    target_retention: Optional[float] = 0.8,
    threshold: Optional[float] = None,
    embedding_field: str = "__embedding",
    cluster_field: str = "__cluster",
    num_tuning_clusters: int = 3,
    dedup_strategy: Literal["random", "nearest", "furthest"] = "random",
    inplace=True,
):
    """Remove semantic near-duplicates from the dataset."""
    if target_retention is None and threshold is None:
        raise ValueError(
            "You must specify either target_retention or threshold."
        )
    if target_retention is not None and threshold is not None:
        logger.warning(
            f"Both target_retention and threshold specified. Using target_retention to tune threshold on {num_tuning_clusters} clusters."
        )

    # first get all the clusters so we only have to do this once. but don't semdedup noise (-1)
    clusters = self._get_clusters(cluster_field)
    if -1 in clusters:
        del clusters[-1]

    if target_retention is not None:
        tuning_clusters = random.choices(
            list(set(clusters.keys())), k=num_tuning_clusters
        )
        # tune threshold
        logger.info(f"Tuning threshold on {num_tuning_clusters} clusters...")
        thresholds = []
        for tuning_cluster in tuning_clusters:
            cluster = clusters[tuning_cluster]
            threshold = tune_threshold(
                cluster,
                self.cluster_centers[cluster_field][tuning_cluster],
                target_retention,
                embedding_field=embedding_field,
                dedup_strategy=dedup_strategy,
            )
            thresholds.append(threshold)
        threshold = np.mean(thresholds)
        logger.info(f"Threshold: {round(threshold, 2)}")

    # get duplicates
    remove = []
    for cluster_id, cluster in clusters.items():
        if (
            len(cluster) < 2 or cluster_id < 0
        ):  # don't deduplicate noise or singletons
            continue
        duplicates = get_duplicates(
            cluster,
            self.cluster_centers[cluster_field][cluster_id],
            threshold=threshold,
            embedding_field=embedding_field,
            dedup_strategy=dedup_strategy,
        )
        logger.info(
            f"Cluster {cluster_id} has {len(duplicates)} duplicates ({round(len(duplicates) / len(cluster) * 100, 1)}%).\n"
        )
        remove.extend(duplicates)

    if inplace:
        before_dedup = len(self.dataset)
        self.dataset = self.dataset.filter(lambda x: x["__id"] not in remove)
        n_removed = len(remove)
        logger.info(
            f"Removed {len(remove)} / {before_dedup} items flagged as semantic near-duplicates ({round(n_removed / before_dedup * 100, 2)}%)."
        )
        return self
    else:
        before_dedup = len(self.dataset)
        new_dataset = self.dataset.filter(lambda x: x["__id"] not in remove)
        n_removed = len(remove)
        logger.info(
            f"Removed {len(remove)} / {before_dedup} items flagged as semantic near-duplicates ({round(n_removed / before_dedup * 100, 2)}%)."
        )
        return type(self)(
            new_dataset,
            model=self.model,
            emb_matrix=self.emb_matrix,
            cluster_ids=self.cluster_ids,
            cluster_centers=self.cluster_centers,
            openai_api_key=self.openai_api_key,
        )


def remove_prototypes(
    self,
    target_retention: float = 0.75,
    cluster_field: str = "__cluster",
    embedding_field: str = "__embedding",
    inplace: bool = True,
):
    """
    Apply the 'SSL prototypes' strategy from D4 (https://arxiv.org/pdf/2308.12284.pdf), which
    removes the most prototypical examples from each cluster (i.e. those closest to the center)
    until the target retention rate is reached.
    """
    # first get all the clusters so we only have to do this once. but don't semdedup noise (-1)
    clusters = self._get_clusters(cluster_field)
    if -1 in clusters:
        del clusters[-1]

    # then keep a running list of what to remove
    remove = []
    for cluster_id, cluster in clusters.items():
        emb_matrix = np.array(cluster[embedding_field])
        cluster_center = self.cluster_centers[cluster_field][cluster_id]
        if (
            len(cluster) < 2 or cluster_id < 0
        ):  # don't deduplicate noise or singletons
            continue
        # sort by distance to center
        dists = np.dot(emb_matrix, cluster_center)
        sorted_indices = np.argsort(dists)
        # remove the top N prototypical examples
        n_to_remove = int(len(cluster) * (1 - target_retention))
        remove.extend(cluster[sorted_indices[:n_to_remove]]["__id"])
        logger.info(
            f"Cluster {cluster_id} has {len(remove)} prototypical examples ({round(len(remove) / len(cluster) * 100, 1)}%).\n"
        )

    if inplace:
        before_dedup = len(self.dataset)
        self.dataset = self.dataset.filter(lambda x: x["__id"] not in remove)
        n_removed = len(remove)
        logger.info(
            f"Removed {len(remove)} / {before_dedup} prototypical items ({round(n_removed / before_dedup * 100, 2)}%)."
        )
        return self
    else:
        before_dedup = len(self.dataset)
        new_dataset = self.dataset.filter(lambda x: x["__id"] not in remove)
        n_removed = len(remove)
        logger.info(
            f"Removed {len(remove)} / {before_dedup} prototypical items ({round(n_removed / before_dedup * 100, 2)}%)."
        )
        return type(self)(
            new_dataset,
            model=self.model,
            emb_matrix=self.emb_matrix,
            cluster_ids=self.cluster_ids,
            cluster_centers=self.cluster_centers,
            openai_api_key=self.openai_api_key,
        )

src/galactic/clustering.py
taylorai/galactic
import re
import os
from typing import Optional, Union, Literal
from .utils import MessageConfig, OPENAI_MESSAGE_CONFIG


def conversation_from_dicts(
    self,
    input_column: str,
    input_message_config: Union[MessageConfig, dict],
    output_column: Optional[str] = None,
    output_message_config: MessageConfig = OPENAI_MESSAGE_CONFIG,
    conversation_field: Optional[str] = None,
    metadata_fields: Optional[list[str]] = [],
):
    """
    Parses (somewhat) arbitrary JSON containing conversations into standard format for OpenAI/etc. conversational fine-tuning.
    If output column is not provided, overwrites the input column.
    """
    if output_column is None:
        output_column = input_column
    if isinstance(input_message_config, dict):
        input_message_config = MessageConfig(**input_message_config)
    role_dict = {
        input_message_config.user_role: output_message_config.user_role,
        input_message_config.assistant_role: output_message_config.assistant_role,
        input_message_config.system_role: output_message_config.system_role,
    }

    def _parse_messages(sample: dict):
        conversation = (
            sample[input_column]
            if conversation_field is None
            else sample[input_column][conversation_field]
        )
        conversation = list(
            map(
                lambda msg: {
                    output_message_config.role_field: role_dict[
                        msg[input_message_config.role_field]
                    ],
                    output_message_config.content_field: msg[
                        input_message_config.content_field
                    ],
                    "metadata": {k: msg[k] for k in metadata_fields},
                },
                conversation,
            )
        )
        return {output_column: conversation}

    self.dataset = self.dataset.map(_parse_messages)


def conversation_from_string(
    self,
    input_column: str,
    user_delimiter: str,
    assistant_delimiter: str,
    system_delimiter: str,
    output_column: Optional[str] = None,
    output_message_config: MessageConfig = OPENAI_MESSAGE_CONFIG,
    strip_whitespace: bool = True,
    replace_values: Optional[dict[str, str]] = None,
):
    """
    Parses a conversation from a string 'transcript' into a structured object with a list of messages.
    """
    if output_column is None:
        output_column = input_column

    pattern = f"({re.escape(user_delimiter)}|{re.escape(assistant_delimiter)}|{re.escape(system_delimiter)})"

    def _parse_single(s: str):
        tokens = re.split(pattern, s)[1:]  # Skip the initial empty string
        result = []
        for i in range(0, len(tokens), 2):
            delimiter, content = tokens[i], tokens[i + 1]
            role = {
                user_delimiter: output_message_config.user_role,
                assistant_delimiter: output_message_config.assistant_role,
                system_delimiter: output_message_config.system_role,
            }.get(delimiter, "unknown")

            if replace_values is not None:
                for k, v in replace_values.items():
                    content = content.replace(k, v)
            if strip_whitespace:
                content = content.strip()

            result.append({"content": content, "role": role})

        return result

    self.dataset = self.dataset.map(
        lambda x: {output_column: _parse_single(x[input_column])}
    )
    return self


def convert_conversation_to_string(
    self,
    input_column: str,
    input_message_config: Union[MessageConfig, dict],
    output_column: str,
    user_delimiter: str,
    assistant_delimiter: str,
    system_delimiter: str,
    strip_whitespace: bool = True,
):
    """
    Converts a conversation from the standard format to a string, using the provided delimiters (e.g. User:, Assistant:)
    """
    if isinstance(input_message_config, dict):
        input_message_config = MessageConfig(**input_message_config)

    def _convert_single(conversation: list[dict]):
        result = ""
        for msg in conversation:
            role, content = (
                msg[input_message_config.role_field],
                msg[input_message_config.content_field],
            )
            result += {
                input_message_config.user_role: user_delimiter,
                input_message_config.assistant_role: assistant_delimiter,
                input_message_config.system_role: system_delimiter,
            }.get(role, "unknown") + content
        if strip_whitespace:
            result = result.strip()
        return result

    self.dataset = self.dataset.map(
        lambda x: {output_column: _convert_single(x[input_column])}
    )
    return self


def get_conversation_length(self, input_column: str, output_column: str):
    """
    Get the length (in messages) of a conversation. If you want length in tokens, convert it to a string, and then use count_tokens
    from the taggers module.
    """
    self.dataset.map(lambda x: {output_column: len(x[input_column])})
    return self


def get_conversation_speakers(
    self,
    input_column: str,
    input_message_config: Union[MessageConfig, dict],
    output_column: str = "__speakers",
):
    """
    Computes list of just the speakers for each message in each conversation.
    """
    self.dataset.map(
        lambda x: {
            output_column: [
                msg[input_message_config.role_field] for msg in x[input_column]
            ]
        }
    )
    return self


def is_alternating(
    self,
    input_column: str,
    input_message_config: Union[MessageConfig, dict],
    output_column: str = "__alternating",
):
    """
    Returns true if no speaker goes 2x in a row. Useful if you make assumptions about conversations being alternating
    and want to catch ones where they're not.
    """

    def _check_alternating(conversation: list[dict]):
        speakers = [
            msg[input_message_config.role_field] for msg in conversation
        ]
        return all(
            speakers[i] != speakers[i + 1] for i in range(len(speakers) - 1)
        )

    self.dataset = self.dataset.map(
        lambda x: {output_column: _check_alternating(x[input_column])}
    )
    return self


def get_last_speaker(
    self,
    input_column: str,
    input_message_config: Union[MessageConfig, dict],
    output_column: str = "__last_speaker",
):
    """
    Gets the role of the last speaker for each conversation.
    """
    if isinstance(input_message_config, dict):
        input_message_config = MessageConfig(**input_message_config)
    self.dataset = self.dataset.map(
        lambda x: {
            output_column: x[input_column][-1][input_message_config.role_field]
        }
    )
    return self


def standardize_last_turn(
    self,
    input_column: str,
    input_message_config: Union[MessageConfig, dict],
    last_speaker_role: Literal["user", "assistant", "system"],
    output_column: Optional[str] = None,
):
    """
    Strip messages from the end of the conversation until the last message is from the specified speaker.
    This is often helpful for training in the case where you want the last reply to be from the assistant,
    and you want to train a model to generate that message given the conversation so far.
    """
    if isinstance(input_message_config, dict):
        input_message_config = MessageConfig(**input_message_config)

    if output_column is None:
        output_column = input_column

    def _standardize_single(conversation: list[dict]):
        expected_role = {
            "user": input_message_config.user_role,
            "assistant": input_message_config.assistant_role,
            "system": input_message_config.system_role,
        }[last_speaker_role]

        while (
            conversation[-1][input_message_config.role_field] != expected_role
        ):
            conversation = conversation[:-1]
        return conversation

    self.dataset = self.dataset.map(
        lambda x: {output_column: _standardize_single(x[input_column])}
    )
    return self


def get_shared_prefix(
    self,
    col1: str,
    col2: str,
    prefix_col: str,
    suffix1_col: Optional[str] = None,
    suffix2_col: Optional[str] = None,
    strip_whitespace: bool = True,
):
    """
    When you have "chosen" and "rejected" conversations, this method will find the longest shared prefix between them.
    Operates on strings, not lists of messages. Resulting prefix will be put in 'prefix_col'. If provided, each suffix
    will be put in the corresponding suffix columns.
    """

    def _get_prefix_single(sample: dict):
        s1, s2 = sample[col1], sample[col2]
        if strip_whitespace:
            s1, s2 = s1.strip(), s2.strip()
        prefix = os.path.commonprefix([s1, s2])
        suffix1 = s1[len(prefix) :]
        suffix2 = s2[len(prefix) :]
        result = {prefix_col: prefix}
        if suffix1_col is not None:
            result[suffix1_col] = suffix1
        if suffix2_col is not None:
            result[suffix2_col] = suffix2
        return result

    self.dataset = self.dataset.map(_get_prefix_single)
    return self


def add_initial_system_message(
    self,
    input_column: str,
    input_message_config: Union[MessageConfig, dict],
    system_message: str,
    output_column: Optional[str] = None,
):
    """
    Adds (the same) initial system message to every conversation.
    """
    if isinstance(input_message_config, dict):
        input_message_config = MessageConfig(**input_message_config)
    if output_column is None:
        output_column = input_column

    new_message = [
        {
            input_message_config.role_field: input_message_config.system_role,
            input_message_config.content_field: system_message,
        }
    ]
    self.dataset.map(lambda x: {output_column: new_message + x[input_column]})
    return self


def take_initial_system_message(
    self,
    input_column: str,
    input_message_config: Union[MessageConfig, dict],
    output_column: str = "__initial_system_message",
    output_type: Literal["dict", "string"] = "string",
    remove_from_input: bool = True,
):
    """
    Takes initial message and puts it in another column if it's a system message.
    If there's no initial system message, it will be [None? empty?]
    """

    def _take_system_msg_single(conversation: list[dict]):
        if conversation[0]["role"] == "system":
            return {
                output_column: conversation[0]
                if output_type == "dict"
                else conversation[0][input_message_config.content_field],
                input_column: conversation[1:]
                if remove_from_input
                else conversation,
            }
        else:
            return {output_column: None, input_column: conversation}

    self.dataset = self.dataset.map(
        lambda x: _take_system_msg_single(x[input_column])
    )
    return self


def take_last_message(
    self,
    input_column: str,
    input_message_config: Union[MessageConfig, dict],
    output_column: str,
    output_type: Literal["dict", "string"] = "string",
    remove_from_input: bool = True,
):
    """
    Take the last message from a conversation and put it in a separate column.
    It will be removed from the input if remove_from_input is True.
    """
    if output_column is None:
        output_column = input_column
    if isinstance(input_message_config, dict):
        input_message_config = MessageConfig(**input_message_config)

    def _take_last_single(conversation: list[dict]):
        return {
            output_column: conversation[-1]
            if output_type == "dict"
            else conversation[-1][input_message_config.content_field],
            input_column: conversation[:-1]
            if remove_from_input
            else conversation,
        }

    self.dataset = self.dataset.map(
        lambda x: _take_last_single(x[input_column])
    )
    return self

src/galactic/conversations.py
taylorai/galactic
from typing import Optional
import numpy as np
from typing import Union, Literal
from .embedding_backends.replicate_backend import ReplicateEmbeddingModel
from .embedding_backends.onnx_backend import ONNXEmbeddingModel
from .embedding_backends.ctranslate2_backend import CT2EmbeddingModel
from .embedding_backends.openai_backend import OpenAIEmbeddingModel
import logging

logger = logging.getLogger("galactic")


def initialize_embedding_model(
    self,
    backend: Literal[
        "cpu", "onnx", "gpu", "openai", "replicate", "modal"
    ] = "cpu",
    model: Literal["gte-small", "gte-tiny", None] = "gte-tiny",
    max_requests_per_minute: Optional[int] = None,
    max_tokens_per_minute: Optional[int] = None,
):
    """
    Initializes the embedding model based on the backend.

    :param backend: The backend used for embedding. Options: 'cpu', 'gpu', 'openai'.
    :type backend: str, default='auto'
    :raises ValueError: Raises an exception if the backend is 'openai' but openai_api_key is not set, or if the backend is unknown.

    Example:

    .. code-block:: python

        initialize_embedding_model('cpu')
    """
    if backend in ["cpu", "onnx"]:
        self.model = ONNXEmbeddingModel(model_name=model)
    elif backend in ["gpu", "ctranslate2"]:
        self.model = CT2EmbeddingModel(model_name=model)

    elif backend == "openai":
        if self.openai_api_key is None:
            raise ValueError(
                "You must set openai_api_key before calling get_embeddings() with openai backend."
            )
        max_rpm = max_requests_per_minute or self.max_requests_per_minute
        max_tpm = max_tokens_per_minute or self.max_tokens_per_minute
        self.model = OpenAIEmbeddingModel(
            openai_api_key=self.openai_api_key,
            max_requests_per_minute=max_rpm,
            max_tokens_per_minute=max_tpm,
        )
    elif backend == "replicate":
        self.model = ReplicateEmbeddingModel(
            replicate_api_key=self.replicate_api_key
        )
    elif backend == "modal":
        from .embedding_backends.modal_backend import ModalEmbeddingModel

        self.model = ModalEmbeddingModel(model_name=model)
    else:
        raise ValueError(f"Unknown backend: {backend}")


def get_embeddings(
    self,
    input_field: str,
    embedding_field: str = "__embedding",
    model: Literal["gte-small", "gte-tiny", None] = "gte-tiny",
    backend: Literal[
        "cpu", "gpu", "openai", "replicate", "modal", "onnx", "ctranslate2"
    ] = "cpu",
    normalize: bool = False,
    pad: bool = False,
    split_strategy: Literal["truncate", "greedy", "even"] = "even",
):
    """
    Get embeddings for a field in the dataset.
    :param input_field: The field to get embeddings for.
    :type input_field: str
    :param embedding_field: The field to store the embeddings in.
    :type embedding_field: str, default='__embedding'
    :param model: The model to use for embeddings, if computing locally. Options: 'gte-small', 'gte-tiny', 'bge-micro'.

    .. code-block:: python

        # Replace "field_name1"
        ds.get_embeddings(field="field_name1")

    """
    # make sure the field doesn't exist already
    if embedding_field in self.dataset.column_names:
        raise ValueError(
            f"Field {embedding_field} already exists in dataset. Please choose a different name, or drop the column."
        )

    self.initialize_embedding_model(model=model, backend=backend)

    # if onnx, don't batch, it doesn't help
    if backend in ["cpu", "onnx"]:
        self.dataset = self.dataset.map(
            lambda x: {
                embedding_field: self.model.embed(
                    x[input_field],
                    normalize=normalize,
                    pad=pad,
                    split_strategy=split_strategy,
                ).tolist()
            }
        )
    # if openai, do one big batch and let our throttling handle rate limits
    elif backend == "openai":
        if normalize:
            logger.info("OpenAI embeddings are normalized by default.")
        else:
            logger.warning(
                "OpenAI embeddings are normalized by default, so normalize=False has no effect."
            )
        embs = self.model.embed_batch(
            self.dataset[input_field], normalize=normalize
        )
        self.dataset = self.dataset.add_column(embedding_field, embs)
    else:
        self.dataset = self.dataset.map(
            lambda x: {
                embedding_field: self.model.embed_batch(x[input_field])
            },
            batched=True,
            batch_size=4096,
        )
    logger.info(
        f"Created embeddings on field '{input_field}', stored in '{embedding_field}'."
    )
    return self


def get_nearest_neighbors(
    self,
    query: Union[str, np.ndarray, list[float]],
    k: int = 5,
    embedding_field: str = "__embedding",
):
    """
    Get the nearest neighbors to a query.

    """
    if embedding_field not in self.dataset.column_names:
        raise ValueError(
            "You must call get_embeddings() before calling get_nearest_neighbors(). If your dataset already has an embeddings column and it's not '__embeddings', pass it as the 'embedding_field' argument."
        )

    emb_matrix = np.array(self.dataset[embedding_field])
    if not hasattr(self, "model"):
        self.initialize_embedding_model()
    if isinstance(query, str):
        query = self.model(query)
    scores = np.dot(emb_matrix, query)
    top_k = list(np.argsort(scores)[::-1][:k])
    top_k_items = self.dataset.select(top_k)
    return top_k_items.to_list()


def reduce_embedding_dim(
    self,
    new_column: Optional[str] = None,
    method: Literal[
        "pca", "incremental_pca", "kernel_pca", "svd", "umap"
    ] = "pca",
    n_dims: int = 50,
    embedding_field: str = "__embedding",
    **kwargs,
):
    n_features = len(self.dataset[embedding_field][0])
    if new_column is None:
        new_column = f"{embedding_field}_{method}_{n_dims}"
    if method == "pca":
        X = np.array(self.dataset[embedding_field])
        from sklearn.decomposition import PCA

        pca = PCA(n_components=n_dims)
        X_transformed = pca.fit_transform(X).tolist()
        self.dataset = self.dataset.add_column(new_column, X_transformed)
    elif method == "incremental_pca":
        from sklearn.decomposition import IncrementalPCA

        pca = IncrementalPCA(n_components=n_dims)
        batch_size = kwargs.get("batch_size", 5 * n_features)

        # fit the PCA on batches of the dataset
        def _fit_batch(batch):
            pca.partial_fit(batch[embedding_field])
            return None

        self.dataset.map(
            _fit_batch,
            batched=True,
            batch_size=batch_size,
        )
        # transform the dataset
        self.dataset = self.dataset.map(
            lambda x: {new_column: pca.transform(x[embedding_field]).tolist()},
            batched=True,
            batch_size=batch_size,
        )
    elif method == "kernel_pca":
        X = np.array(self.dataset[embedding_field])
        from sklearn.decomposition import KernelPCA

        pca = KernelPCA(
            n_components=n_dims, kernel=kwargs.get("kernel", "rbf")
        )
        X_transformed = pca.fit_transform(X).tolist()
        self.dataset = self.dataset.add_column(new_column, X_transformed)
    elif method == "svd":
        X = np.array(self.dataset[embedding_field])
        from sklearn.decomposition import TruncatedSVD

        svd = TruncatedSVD(n_components=n_dims)
        X_transformed = svd.fit_transform(X).tolist()
        self.dataset = self.dataset.add_column(new_column, X_transformed)

    elif method == "umap":
        # warn that it will be slow if embedding dim is large
        if len(self.dataset[embedding_field][0]) >= 100:
            logger.info(
                "Embedding dimension is large, which can slow UMAP down a lot, even for small datasets. You might consider PCA first with e.g. n_dims=50, followed by UMAP."
            )
        import umap

        reducer = umap.UMAP(n_components=n_dims, **kwargs)
        X = np.array(self.dataset[embedding_field])
        X_transformed = reducer.fit_transform(X).tolist()
        self.dataset = self.dataset.add_column(new_column, X_transformed)
    else:
        raise ValueError(f"Unknown method: {method}")

    logger.info(
        f"Reduced embeddings to {n_dims} dimensions using {method}. New embeddings stored in column '{new_column}'"
    )

    return self


# finetune embeddings as in https://github.com/openai/openai-cookbook/blob/main/examples/Customizing_embeddings.ipynb
def tune_embeddings(self):
    pass

src/galactic/embedding.py
taylorai/galactic
from abc import ABC, abstractmethod
from typing import Union, Tuple, Optional, Literal
import numpy as np
import logging

logger = logging.getLogger("galactic")


class EmbeddingModelBase(ABC):
    def __init__(self):
        super().__init__()
        self.tokenizer = None
        self.max_length = 512

    def __call__(
        self,
        input: Union[str, list[str]],
    ) -> Union[list[list[float]], list[float], np.ndarray]:
        if isinstance(input, str):
            return self.embed(input)
        elif isinstance(input, list):
            return self.embed_batch(input)
        else:
            raise ValueError(
                f"Input must be str or list[str], not {type(input)}"
            )

    def split_and_tokenize_single(
        self,
        text: str,
        pad: bool = False,
        split_strategy: Literal["truncate", "greedy", "even"] = "even",
    ) -> dict[str, list[list[int]]]:
        """
        Split and tokenize a single text to prepare it for the embedding model.
        Padding is only necessary if running more than 1 sequence thru the model at once.
        Splitting happens when the model exceeds the max_length (usually 512).
        You can either truncate the text, or split into chunks. Chunking can be "greedy"
        (as many 512 chunks as possible), or "even" (split into even-ish chunks with np.array_split).
        """

        # first make into tokens
        tokenized = self.tokenizer(text)  # (seq_len, )

        # if don't have to pad and don't have to split into chunks, we're done
        if not pad and len(tokenized["input_ids"]) <= self.max_length:
            return {k: [tokenized[k]] for k in tokenized}

        # handle splitting
        if split_strategy == "truncate":
            for k in tokenized:
                tokenized[k] = [tokenized[k][: self.max_length]]

        elif split_strategy == "greedy":
            for k in tokenized:
                tokenized[k] = [
                    tokenized[k][idx : idx + self.max_length]
                    for idx in range(0, len(tokenized[k]), self.max_length)
                ]

        elif split_strategy == "even":
            for k in tokenized:
                tokenized[k] = [
                    arr.tolist()
                    for arr in np.array_split(
                        tokenized[k],
                        int(np.ceil(len(tokenized[k]) / self.max_length)),
                    )
                ]

        else:
            raise ValueError(
                f"split_strategy must be 'truncate', 'greedy', or 'even', not {split_strategy}"
            )

        # pad if applicable
        if pad:
            # first make sure list is nested
            if not isinstance(tokenized["input_ids"][0], list):
                for k in tokenized:
                    tokenized[k] = [tokenized[k]]

            # get pad token
            pad_token_id = self.tokenizer.pad_token_id
            if pad_token_id is None:
                pad_token_id = 0

            pad_len = max(
                [
                    len(tokenized["input_ids"][i])
                    for i in range(len(tokenized["input_ids"]))
                ]
            )
            for k in tokenized:
                tokenized[k] = [
                    np.pad(
                        tokenized[k][i],
                        (0, pad_len - len(tokenized[k][i])),
                        constant_values=pad_token_id,
                    ).tolist()
                    for i in range(len(tokenized[k]))
                ]

        return tokenized

    def split_and_tokenize_batch(
        self,
        texts: str,
        pad: bool = False,
        split_strategy: Literal["truncate", "greedy", "even"] = "even",
    ) -> dict:
        """
        Tokenize the text and pad if applicable.

        :param text: The input text to be tokenized.
        :type text: str
        :return: Returns a tuple. dictionary containing tokenized and padded 'input_ids',
        'attention_mask' and 'token_type_ids', along with a list of offsets.
        :rtype: Tuple[Dict[str, numpy.ndarray], list[int]]

        Example:

        .. code-block:: python

            tokenized_text = model.split_and_tokenize('sample text')
        """
        result = {}
        offsets = [0]
        if self.tokenizer is None:
            raise ValueError("Tokenizer is not initialized.")
        if self.max_length is None:
            raise ValueError("max_length is not initialized.")

        # first tokenize without padding
        for text in texts:
            tokenized = self.split_and_tokenize_single(
                text, pad=False, split_strategy=split_strategy
            )
            for k in tokenized:
                if k not in result:
                    result[k] = tokenized[k]
                else:
                    result[k].extend(tokenized[k])

            offsets.append(len(result["input_ids"]))

        # then, if padding, use longest length in batch
        if pad:
            pad_token_id = self.tokenizer.pad_token_id
            if pad_token_id is None:
                pad_token_id = 0

            pad_len = max([len(tokenized[k][0]) for k in result])
            for k in result:
                result[k] = [
                    np.pad(
                        result[k][i],
                        (0, pad_len - len(result[k][i])),
                        constant_values=pad_token_id,
                    ).tolist()
                    for i in range(len(result[k]))
                ]

        return {
            "tokens": result,
            "offsets": offsets,
        }

    @abstractmethod
    def embed(
        self, text: str, normalize: bool = False
    ) -> Union[list[float], np.ndarray]:
        pass

    @abstractmethod
    def embed_batch(
        self, texts: list[str], normalize: bool = False
    ) -> Union[list[list[float]], np.ndarray]:
        pass

src/galactic/embedding_backends/base.py
taylorai/galactic
import os
import multiprocessing
import torch
import ctranslate2
from .base import EmbeddingModelBase
import numpy as np
from transformers import AutoTokenizer
from typing import Literal
from huggingface_hub import snapshot_download
import logging

logger = logging.getLogger("galactic")

model_registry = {
    "gte-small": {
        "remote_file": "gte-small",
        "tokenizer": "Supabase/gte-small",
    },
    "gte-tiny": {
        "remote_file": "gte-tiny",
        "tokenizer": "Supabase/gte-small",
    },
    "bge-micro": {
        "remote_file": "bge-micro",
        "tokenizer": "BAAI/bge-small-en-v1.5",
    },
}


class CT2EmbeddingModel(EmbeddingModelBase):
    def __init__(
        self,
        model_name: str = "bge-micro",
        device: str = "gpu",
        max_length: int = 512,
    ):
        super().__init__()

        # select device
        if device == "gpu":
            if torch.cuda.is_available():
                self.device = "cuda"
            else:
                logger.warning("GPU not available, using CPU.")
                self.device = "cpu"

        # download model
        model_path = snapshot_download(
            "TaylorAI/galactic-models",
            allow_patterns=[model_registry[model_name]["remote_file"] + "/*"],
        )
        model_path = os.path.join(
            model_path, model_registry[model_name]["remote_file"]
        )
        self.model = ctranslate2.Encoder(
            model_path,
            device=self.device,
            intra_threads=multiprocessing.cpu_count(),
        )
        self.max_length = max_length
        self.tokenizer = AutoTokenizer.from_pretrained(
            model_registry[model_name]["tokenizer"]
        )

        # select batch size
        if self.device == "cuda":
            self.batch_size = self._select_batch_size()
        else:
            self.batch_size = 1

    def _select_batch_size(self):
        current_batch_size = 4096
        sample_batch = [[11] * self.max_length] * current_batch_size
        while current_batch_size > 0:
            try:
                self.model.forward_batch(sample_batch)
                logger.info(f"Batch size selected: {current_batch_size}")
                return current_batch_size
            except Exception as e:
                if "out of memory" in str(e):
                    current_batch_size = current_batch_size // 2
                    sample_batch = sample_batch[:current_batch_size]
                else:
                    raise e

        logger.error(
            "Failed with batch size 1. You'll need a bigger GPU or a smaller model."
        )
        raise RuntimeError(
            "Failed with batch size 1. You'll need a bigger GPU or a smaller model."
        )

    def embed(
        self,
        text: str,
        normalize: bool = False,
        pad: bool = False,
        split_strategy: Literal["truncate", "greedy", "even"] = "even",
    ):
        if pad:
            logger.warning(
                "Padding is not necessary for CTranslate2. Ignoring pad=True."
            )

        input = self.split_and_tokenize_single(
            text, pad=False, split_strategy=split_strategy
        )
        output = self.model.forward_batch(input["input_ids"]).pooler_output
        if self.device == "cuda":
            output = (
                torch.as_tensor(output, device="cuda")
                .cpu()
                .numpy()
                .mean(axis=0)
            )
        else:
            output = np.array(output).mean(axis=0)

        if normalize:
            output = output / np.linalg.norm(output)

        return output

    def embed_batch(
        self,
        texts: list[str],
        normalize: bool = False,
        pad: bool = False,
        split_strategy: Literal["truncate", "greedy", "even"] = "even",
    ):
        if pad:
            logger.warning(
                "Padding is not necessary for CTranslate2. Ignoring pad=True."
            )

        tokenized = self.split_and_tokenize_batch(
            texts, pad=False, split_strategy=split_strategy
        )
        inputs = tokenized["tokens"]
        offsets = tokenized["offsets"]
        outputs = None
        for i in range(0, len(inputs["input_ids"]), self.batch_size):
            batch = inputs["input_ids"][i : i + self.batch_size]
            batch_out = self.model.forward_batch(batch).pooler_output
            if self.device == "cuda":
                batch_out = (
                    torch.as_tensor(batch_out, device="cuda")
                    .cpu()
                    .numpy()  # batch_size, hidden_size
                )
            else:
                batch_out = np.array(batch_out)  # batch_size, hidden_size
            if outputs is None:
                outputs = batch_out
            else:
                outputs = np.concatenate([outputs, batch_out], axis=0)

        # use offsets to average each text's embeddings
        embs = []
        for i in range(len(offsets) - 1):
            start, end = offsets[i], offsets[i + 1]
            chunk = outputs[start:end]
            averaged = chunk.mean(axis=0)
            if normalize:
                averaged = averaged / np.linalg.norm(averaged)
            embs.append(averaged)

        return np.array(embs)

src/galactic/embedding_backends/ctranslate2_backend.py
taylorai/galactic
# if you want to use Modal Labs for embeddings, you'll need to deploy this in YOUR modal account under the name 'gte_small'.
# then, if you're authenticated to Modal, you should be able to look up the function and use it as an embeddings backend.
import numpy as np
from modal import Cls
from typing import Union, Literal


class ModalEmbeddingModel:
    def __init__(self, model_name="bge-micro"):
        ModalEmbeddingModel = Cls.lookup(
            "modal_embeddings", "RemoteEmbeddingModel"
        )
        self.model = ModalEmbeddingModel(model_name=model_name)

    def embed(
        self,
        text: str,
        normalize: bool = False,
        pad: bool = False,
        split_strategy: Literal["truncate", "greedy", "even"] = "even",
    ):
        return self.model.embed_batch.remote(
            [text], normalize=normalize, pad=pad, split_strategy=split_strategy
        ).reshape(-1)

    def embed_batch(
        self,
        texts: list[str],
        batch_size: int = 250,
        normalize: bool = False,
        pad: bool = False,
        split_strategy: Literal["truncate", "greedy", "even"] = "even",
    ):
        chunks = [
            texts[i : i + batch_size] for i in range(0, len(texts), batch_size)
        ]
        promises = self.model.embed_batch.map(
            chunks,
            kwargs={
                "normalize": normalize,
                "pad": pad,
                "split_strategy": split_strategy,
            },
        )
        results = [p for p in promises]  # list of np arrays
        flattened = np.concatenate(results, axis=0)
        return flattened

src/galactic/embedding_backends/modal_backend.py
taylorai/galactic
# if you want to use Modal Labs for embeddings, you'll need to deploy this in YOUR modal account under the name 'gte_small'.
# then, if you're authenticated to Modal, you should be able to look up the function and use it as an embeddings backend.
from modal import Image, Stub, method, Cls
from typing import Union, Literal

image = (
    # Python 3.11+ not yet supported for torch.compile
    Image.debian_slim(python_version="3.10")
    .apt_install("git")
    .pip_install(
        "torch",
        "galactic-ai>=0.2.15",
        "huggingface-hub",
        "hf-transfer~=0.1",
    )
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
)

stub = Stub(name="modal_embeddings", image=image)


# Embedding model that runs on Modal, wrapping the ONNX embedding model.
@stub.cls()
class RemoteEmbeddingModel:
    def __enter__(self):
        from huggingface_hub import hf_hub_download

        model_registry = {
            "gte-small": {
                "remote_file": "model_quantized.onnx",
                "tokenizer": "Supabase/gte-small",
            },
            "gte-tiny": {
                "remote_file": "gte-tiny.onnx",
                "tokenizer": "Supabase/gte-small",
            },
            "bge-micro": {
                "remote_file": "bge-micro.onnx",
                "tokenizer": "BAAI/bge-small-en-v1.5",
            },
        }

        # download allll the things!
        for model_name in model_registry:
            hf_hub_download(
                "TaylorAI/galactic-models",
                filename=model_registry[model_name]["remote_file"],
            )

    def __init__(self, model_name="bge-micro"):
        from galactic.embedding_backends.onnx_backend import ONNXEmbeddingModel

        self.model = ONNXEmbeddingModel(
            model_name=model_name,
        )

    @method()
    def embed(
        self,
        text: str,
        normalize: bool = False,
        pad: bool = False,
        split_strategy: Literal["truncate", "greedy", "even"] = "even",
    ):
        return self.model.embed(
            text, normalize=normalize, pad=pad, split_strategy=split_strategy
        )

    @method()
    def embed_batch(
        self,
        texts: list[str],
        normalize: bool = False,
        pad: bool = False,
        split_strategy: Literal["truncate", "greedy", "even"] = "even",
    ):
        return self.model.embed_batch(
            texts, normalize=normalize, pad=pad, split_strategy=split_strategy
        )

src/galactic/embedding_backends/modal_remote.py
taylorai/galactic
from .base import EmbeddingModelBase
import onnxruntime as ort
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download
import numpy as np
from typing import Literal

import logging

logger = logging.getLogger("galactic")


class ONNXEmbeddingModel(EmbeddingModelBase):
    def __init__(
        self,
        model_name: str = "bge-micro",
        max_length: int = 512,
        local_path=None,
    ):
        super().__init__()
        self.model_registry = {
            "gte-small": {
                "remote_file": "model_quantized.onnx",
                "tokenizer": "Supabase/gte-small",
            },
            "gte-tiny": {
                "remote_file": "gte-tiny.onnx",
                "tokenizer": "Supabase/gte-small",
            },
            "bge-micro": {
                "remote_file": "bge-micro.onnx",
                "tokenizer": "BAAI/bge-small-en-v1.5",
            },
        }
        if local_path is None:
            local_path = hf_hub_download(
                "TaylorAI/galactic-models",
                filename=self.model_registry[model_name]["remote_file"],
            )
        self.tokenizer = AutoTokenizer.from_pretrained(
            self.model_registry[model_name]["tokenizer"]
        )
        self.max_length = max_length
        self.providers = ["CPUExecutionProvider"]
        # # this only matters if using onnxruntime-silicon which didn't do anything for me
        # if "CoreMLExecutionProvider" in ort.get_available_providers():
        #     self.providers.append("CoreMLExecutionProvider")
        # if "CUDAExecutionProvider" in ort.get_available_providers():
        #     logger.info("Using CUDAExecutionProvider since it is available.")
        #     self.providers = ["CUDAExecutionProvider"]
        #     if "TensorrtExecutionProvider" in ort.get_available_providers():
        #         logger.info(
        #             "Using TensorrtExecutionProvider since it is available."
        #         )
        #         self.providers.append("TensorrtExecutionProvider")
        self.session = ort.InferenceSession(
            local_path, providers=self.providers
        )

    def embed(
        self,
        text: str,
        normalize: bool = False,
        pad: bool = False,
        split_strategy: Literal["truncate", "greedy", "even"] = "even",
    ):
        input = self.split_and_tokenize_single(
            text, pad=pad, split_strategy=split_strategy
        )
        outs = []
        for seq in range(len(input["input_ids"])):
            if not pad:
                assert (
                    np.mean(input["attention_mask"][seq]) == 1
                ), "pad=False but attention_mask has 0s"
            out = self.session.run(
                None,
                {
                    "input_ids": input["input_ids"][seq : seq + 1],
                    "attention_mask": input["attention_mask"][seq : seq + 1],
                    "token_type_ids": input["token_type_ids"][seq : seq + 1],
                },
            )[
                0
            ]  # 1, seq_len, hidden_size
            trimmed = out[
                0, np.array(input["attention_mask"][seq]) == 1, :
            ]  # chunk_seq_len, hidden_size
            outs.append(trimmed)
        outs = np.concatenate(outs, axis=0)  # full_seq_len, hidden_size
        avg = np.mean(outs, axis=0)  # hidden_size
        if normalize:
            avg = avg / np.linalg.norm(avg)
        return avg

    def embed_batch(
        self,
        texts: list[str],
        normalize: bool = False,
        pad: bool = False,
        split_strategy: Literal["truncate", "greedy", "even"] = "even",
    ):
        result = []
        for text in texts:
            result.append(
                self.embed(
                    text,
                    normalize=normalize,
                    pad=pad,
                    split_strategy=split_strategy,
                )
            )

        return np.array(result)

src/galactic/embedding_backends/onnx_backend.py
taylorai/galactic
from .base import EmbeddingModelBase
from ..async_openai import embed_texts_with_openai
import tiktoken
import openai
import numpy as np


class OpenAIEmbeddingModel(EmbeddingModelBase):
    def __init__(
        self,
        openai_api_key: str,
        max_requests_per_minute: int,
        max_tokens_per_minute: int,
    ):
        super().__init__()
        self.tokenizer = tiktoken.get_encoding("cl100k_base")
        self.openai_api_key = openai_api_key
        self.max_requests_per_minute = max_requests_per_minute
        self.max_tokens_per_minute = max_tokens_per_minute

    def embed(self, text: str, normalize: bool = False):
        """
        Embeds a single text with OpenAI API.

        :param text: The input text to be embedded.
        :type text: str
        :return: The normalized averaged embeddings.
        :rtype: numpy.ndarray
        """
        tokens = self.tokenizer.encode(text)
        # if longer than 8191, split into chunks
        if len(tokens) > 8191:
            num_chunks = int(np.ceil(len(tokens) / 8191))
            chunks = np.array_split(tokens, num_chunks).tolist()
        else:
            chunks = [tokens]
        res = openai.Embedding.create(
            model="text-embedding-ada-002",
            input=chunks,
        )
        embs = np.array([res.data[i].embedding for i in range(len(res.data))])
        avg = np.mean(embs, axis=0)
        return avg / np.linalg.norm(avg)

    def embed_batch(self, texts: list[str], normalize: bool = False):
        """
        Embeds a batch of texts with OpenAI API.
        :param texts: The input texts to be embedded.
        :type texts: list[str]
        :return: The normalized averaged embeddings.

        """
        embs = embed_texts_with_openai(
            texts,
            self.openai_api_key,
            max_tokens_per_minute=self.max_tokens_per_minute,
            max_requests_per_minute=self.max_requests_per_minute,
            show_progress=True,
        )
        return embs

src/galactic/embedding_backends/openai_backend.py
taylorai/galactic
### Utils for making lots of requests to the Replicate API & throttling to avoid hitting rate limit.
### NOT CURRENTLY USED IN PRODUCTION, here for reference & possible future use for embeddings backend.
### Code here adapted from openai cookbook: https://github.com/openai/openai-cookbook/blob/main/examples/api_request_parallel_processor.py
import aiohttp
import asyncio
import json
import logging
import time
import random
from tqdm.auto import tqdm
from dataclasses import dataclass, field
from itertools import islice
from .base import EmbeddingModelBase

logger = logging.getLogger("galactic")


class ReplicateEmbeddingModel(EmbeddingModelBase):
    def __init__(
        self,
        replicate_api_key: str,
        embedding_url="https://api.replicate.ai/v1/predictions",
        api_version="3358c35efbe0c1ca88f6ed77dcd9d306fff9ddf0341f51747ea10ee494ddd682",
        max_requests_per_minute: int = 3000,
    ):
        super().__init__()
        self.replicate_api_key = replicate_api_key
        self.embedding_url = embedding_url
        self.api_version = api_version
        self.max_requests_per_minute = max_requests_per_minute

    # should probably just do this with onnx for a single embedding??
    def embed(self, text: str):
        emb = embed_texts_with_replicate(
            texts=[text],
            api_key=self.replicate_api_key,
            url=self.embedding_url,
            api_version=self.api_version,
            max_requests_per_minute=self.max_requests_per_minute,
        )[0]
        return emb

    def embed_batch(self, texts: list[str]):
        embs = embed_texts_with_replicate(
            texts=texts,
            api_key=self.replicate_api_key,
            url=self.embedding_url,
            api_version=self.api_version,
            max_requests_per_minute=self.max_requests_per_minute,
        )
        return embs


@dataclass
class StatusTracker:
    num_tasks_started: int = 0
    num_tasks_in_progress: int = 0
    num_tasks_succeeded: int = 0
    num_tasks_failed: int = 0
    num_rate_limit_errors: int = 0
    time_of_last_rate_limit_error: int = 0
    total_requests = 0


@dataclass
class APIRequest:
    task_id: int
    texts: list[str]
    attempts_left: int
    url: str
    api_version: str
    result: list = field(default_factory=list)

    def __post_init__(self):
        # get the URL and request JSON
        self.request_json = {
            "version": self.api_version,
            "input": {
                "texts": json.dumps(self.texts),
            },
        }

    async def call_api(
        self,
        request_header: dict,
        retry_queue: asyncio.Queue,
        status_tracker: StatusTracker,
        pbar: tqdm,
    ):
        try:
            status_tracker.total_requests += 1
            timeout = aiohttp.ClientTimeout(total=20)
            async with aiohttp.ClientSession(timeout=timeout) as session:
                # first create the prediction
                async with session.post(
                    url=self.url,
                    headers=request_header,
                    json=self.request_json,
                ) as response:
                    res_json = await response.json()

                # if status is 201, prediction was successfully created
                if response.status == 201:
                    # poll the prediction until it's ready. exponential backoff up to 1 minute. 5 -> 10 -> 20 -> 40 -> 60
                    # don't wait for more than 5 minutes.
                    prediction_url = res_json["urls"]["get"]
                    # random jitter to avoid thundering herd
                    time_to_wait = 5
                    res_json = None
                    while time_to_wait < 500:
                        await asyncio.sleep(
                            time_to_wait + random.random() * time_to_wait / 2
                        )
                        async with session.get(
                            url=prediction_url,
                            headers=request_header,
                        ) as response:
                            res_json = await response.json()

                        if res_json["status"] == "succeeded":
                            pbar.update(1)
                            self.result.append(res_json)
                            status_tracker.num_tasks_in_progress -= 1
                            status_tracker.num_tasks_succeeded += 1
                            return
                        elif res_json["status"] in ["failed", "canceled"]:
                            logger.warning("Prediction failed or canceled.")
                            self.result.append(res_json)
                            break
                        elif res_json["status"] in ["starting", "processing"]:
                            # still processing the prediction, poll again later
                            time_to_wait += 15
                        elif response.status == 429:
                            logger.warning(
                                "Rate limit error while fetching prediction."
                            )
                            time_to_wait += 15
                    if res_json is None:
                        logger.error("Prediction polling timed out.")

                elif response.status == 429:
                    status_tracker.time_of_last_rate_limit_error = time.time()
                    status_tracker.num_rate_limit_errors += 1
                    logger.warning("Rate limit error.")
                    self.result.append(res_json)
                else:
                    logger.warning("Other error. " + str(res_json))
                    self.result.append(res_json)

                if self.attempts_left:
                    logger.info(f"Will retry request {self.task_id}.")
                    self.attempts_left -= 1
                    retry_queue.put_nowait(self)
                else:
                    print("out of tries")
                    status_tracker.num_tasks_in_progress -= 1
                    status_tracker.num_tasks_failed += 1

        except Exception as e:
            self.result.append(str(e))
            if self.attempts_left:
                retry_queue.put_nowait(self)
            else:
                status_tracker.num_tasks_in_progress -= 1
                status_tracker.num_tasks_failed += 1


async def process_api_requests_from_list(
    texts: list[list[str]],
    api_key: str,
    max_attempts: int,
    url: str,
    api_version: str,
    max_requests_per_minute: int = 3000,
):
    """Processes API requests in parallel, throttling to stay under rate limits."""
    # constants
    seconds_to_pause_after_rate_limit_error = 15
    seconds_to_sleep_each_loop = (
        0.01  # 10 ms limits max throughput to 100 requests per second
    )

    # infer API endpoint and construct request header
    request_header = {"Authorization": f"Token {api_key}"}

    # initialize trackers
    queue_of_requests_to_retry = asyncio.Queue()
    status_tracker = (
        StatusTracker()
    )  # single instance to track a collection of variables
    next_request = None  # variable to hold the next request to call

    # initialize available capacity counts
    available_request_capacity = max_requests_per_minute
    last_update_time = time.time()

    # initialize flags
    texts_not_finished = True
    logger.debug(f"Initialization complete.")

    # turn the texts into an iterator
    pbar = tqdm(total=len(texts))
    texts = iter(enumerate(texts))
    results = []
    while True:
        # get next request (if one is not already waiting for capacity)
        if next_request is None:
            if not queue_of_requests_to_retry.empty():
                next_request = queue_of_requests_to_retry.get_nowait()
                logger.debug(
                    f"Retrying request {next_request.task_id}: {next_request}"
                )
            elif texts_not_finished:
                try:
                    # get new request
                    idx, text_batch = next(texts)
                    next_request = APIRequest(
                        task_id=idx,
                        texts=text_batch,
                        attempts_left=max_attempts,
                        url=url,
                        api_version=api_version,
                    )
                    status_tracker.num_tasks_started += 1
                    status_tracker.num_tasks_in_progress += 1
                    logger.debug(
                        f"Reading request {next_request.task_id}: {next_request}"
                    )
                except StopIteration:
                    # if file runs out, set flag to stop reading it
                    logger.debug("Read file exhausted")
                    texts_not_finished = False

        # update available capacity
        current_time = time.time()
        seconds_since_update = current_time - last_update_time
        available_request_capacity = min(
            available_request_capacity
            + max_requests_per_minute * seconds_since_update / 60.0,
            max_requests_per_minute,
        )
        last_update_time = current_time

        # if enough capacity available, call API
        if next_request:
            if available_request_capacity >= 1:
                # update counters
                available_request_capacity -= 1
                next_request.attempts_left -= 1

                # call API
                asyncio.create_task(
                    next_request.call_api(
                        request_header=request_header,
                        retry_queue=queue_of_requests_to_retry,
                        status_tracker=status_tracker,
                        pbar=pbar,
                    )
                )
                logger.debug(
                    f"Called API for request {next_request.task_id}: {next_request}"
                )
                results.append(next_request)
                next_request = None  # reset next_request to empty

        # if all tasks are finished, break
        if status_tracker.num_tasks_in_progress == 0:
            break

        # main loop sleeps briefly so concurrent tasks can run
        await asyncio.sleep(seconds_to_sleep_each_loop)

        # if a rate limit error was hit recently, pause to cool down
        seconds_since_rate_limit_error = (
            time.time() - status_tracker.time_of_last_rate_limit_error
        )
        if (
            seconds_since_rate_limit_error
            < seconds_to_pause_after_rate_limit_error
        ):
            remaining_seconds_to_pause = (
                seconds_to_pause_after_rate_limit_error
                - seconds_since_rate_limit_error
            )
            await asyncio.sleep(remaining_seconds_to_pause)
            # ^e.g., if pause is 15 seconds and final limit was hit 5 seconds ago
            logger.warn(
                f"Pausing to cool down until {time.ctime(status_tracker.time_of_last_rate_limit_error + seconds_to_pause_after_rate_limit_error)}"
            )

    # after finishing, log final status
    logger.info(f"""Parallel processing complete.""")
    if status_tracker.num_tasks_failed > 0:
        logger.warning(
            f"{status_tracker.num_tasks_failed} / {status_tracker.num_tasks_started} requests failed."
        )
    if status_tracker.num_rate_limit_errors > 0:
        logger.warning(
            f"{status_tracker.num_rate_limit_errors} rate limit errors received. Consider running at a lower rate."
        )
    return results


def batchify(data, batch_size):
    iter_data = iter(data)
    while True:
        batch = list(islice(iter_data, batch_size))
        if not batch:
            break
        yield batch


def embed_texts_with_replicate(
    texts: list[str],
    api_key: str,
    url: str,
    api_version: str,
    batch_size: int = 25,
    max_attempts: int = 3,
    max_requests_per_minute: int = 3000,
):
    results = asyncio.run(
        process_api_requests_from_list(
            texts=list(batchify(texts, batch_size)),
            api_key=api_key,
            max_attempts=max_attempts,
            max_requests_per_minute=max_requests_per_minute,
            url=url,
            api_version=api_version,
        )
    )
    # extract the embeddings
    embs = []
    for result in sorted(results, key=lambda x: x.task_id):
        embs.extend(result.result[-1]["output"])
    print("Number of texts:", len(texts), "Number of embs:", len(embs))
    return embs

src/galactic/embedding_backends/replicate_backend.py
taylorai/galactic
import platform
from tempfile import TemporaryDirectory
from pathlib import Path
import pytesseract
from pdf2image import convert_from_path
from PIL import Image
import logging

logger = logging.getLogger("galactic")

if platform.system() == "Windows":
    pytesseract.pytesseract.tesseract_cmd = (
        r"C:\Program Files\Tesseract-OCR\tesseract.exe"
    )
    path_to_poppler_exe = Path(r"C:\.....")
    out_directory = Path(r"~\Desktop").expanduser()
else:
    out_directory = Path("~").expanduser()


def extract_doc(path: str) -> str:
    image_file_list = []
    extracted_text = ""  # This variable will accumulate the extracted text

    with TemporaryDirectory() as tempdir:
        if platform.system() == "Windows":
            pdf_pages = convert_from_path(
                path, 500, poppler_path=path_to_poppler_exe
            )
        else:
            pdf_pages = convert_from_path(path, 500)

        for page_enumeration, page in enumerate(pdf_pages, start=1):
            filename = Path(tempdir) / f"page_{page_enumeration:03}.jpg"
            page.save(filename, "JPEG")
            image_file_list.append(filename)

        for image_file in image_file_list:
            with Image.open(image_file) as img:
                text = str(pytesseract.image_to_string(img))
                text = text.replace("-\n", "")
                extracted_text += text + "\n"  # Accumulate the text

    return extracted_text  # Return the accumulated text


"""
from src.galactic import GalacticDataset
from src.galactic import extract_doc
extract_doc("../Downloads/example.pdf")
"""

src/galactic/extract_doc.py
taylorai/galactic
# there are 2 types of filter: stateless filter that can operate on only a single example,
# or a stateful filter (e.g. Bloom filter) where examples depend on each other.
import re
import json
import pybloom_live
from typing import Sequence

import logging

logger = logging.getLogger("galactic")


def apply_bloom_filter(self, fields: Sequence[str], inplace: bool = True):
    """
    Applies a Bloom filter to the dataset to filter out duplicate examples based on the specified fields.

    .. code-block:: python

        # Example usage:
        ds.apply_bloom_filter(fields=['field1', 'field2'], inplace=False)

    :param fields: Sequence of field names used to apply the Bloom filter.
    :param inplace: Whether to apply the filter in-place or return a new object. Defaults to True.
    :return: If inplace is False, returns a new object of the same type with the filtered dataset.
    :rtype: type(self) if not inplace else None
    """

    bloom = pybloom_live.BloomFilter(
        capacity=len(self.dataset)
        if len(self.dataset) < int(1.0e9)
        else int(1.0e9),
        error_rate=0.001,
    )

    def bloom_filter(sample):
        key = json.dumps({field: sample[field] for field in fields})
        if key in bloom:
            return False
        else:
            bloom.add(key)
            return True

    if inplace:
        self.dataset = self.dataset.filter(bloom_filter)
        logger.info(
            f"Filtered dataset in-place with Bloom filter on fields: {fields}"
        )
        # return self for chaining
        return self
    else:
        new_dataset = self.dataset.filter(bloom_filter)
        logger.info(f"Filtered dataset with Bloom filter on fields: {fields}")
        return type(self)(
            new_dataset,
            model=self.model,
            emb_matrix=self.emb_matrix,
            cluster_ids=self.cluster_ids,
            cluster_centers=self.cluster_centers,
            openai_api_key=self.openai_api_key,
        )


# these filters are methods to attach to the GalacticDataset class.
# each filter takes a list of fields which the filter will be (separately) applied to.
def filter_string(
    self, fields: Sequence[str], values: Sequence[str], inplace: bool = True
):
    """
    Filters the dataset by removing examples that contain any of the specified strings in the specified fields.

    .. code-block:: python

        # Example usage:
        ds.filter_string(fields=['field1', 'field2'], values=['value1', 'value2'], inplace=False)

    :param fields: Sequence of field names to apply the filter on.
    :param values: Sequence of string values to filter out from the specified fields.
    :param inplace: Whether to apply the filter in-place or return a new object. Defaults to True.
    :return: If inplace is False, returns a new object of the same type with the filtered dataset.
    :rtype: type(self) if not inplace else None
    """

    # make a regexp that matches any of the fields
    regexp = re.compile("|".join([re.escape(val) for val in values]))

    def filter_(sample):
        for field in fields:
            if field in sample:
                if isinstance(sample[field], str):
                    if regexp.search(sample[field]):
                        return False
                else:
                    if regexp.search(str(sample[field])):
                        return False
        return True

    if inplace:
        self.dataset = self.dataset.filter(filter_)
        logger.info(
            f"Filtered dataset in-place with exact string matching on fields: {fields}"
        )
        # return self for chaining
        return self
    else:
        new_dataset = self.dataset.filter(filter_)
        logger.info(
            f"Filtered dataset with exact string matching on fields: {fields}"
        )
        return type(self)(
            new_dataset,
            model=self.model,
            emb_matrix=self.emb_matrix,
            cluster_ids=self.cluster_ids,
            cluster_centers=self.cluster_centers,
            openai_api_key=self.openai_api_key,
        )


def filter_regex(
    self, fields: Sequence[str], regex: str, inplace: bool = True
):
    """
    Filters the dataset by removing examples that match the specified regex pattern in the specified fields.

    .. code-block:: python

        # Example usage:
        ds.filter_regex(fields=['field1', 'field2'], regex='[0-9]+', inplace=False)

    :param fields: Sequence of field names to apply the regex filter on.
    :param regex: The regex pattern used to filter out matching examples from the specified fields.
    :param inplace: Whether to apply the filter in-place or return a new object. Defaults to True.
    :return: If inplace is False, returns a new object of the same type with the filtered dataset.
    :rtype: type(self) if not inplace else None
    """
    # make a regexp that matches any of the fields
    regexp = re.compile(regex)

    def filter_(sample):
        for field in fields:
            if field in sample:
                if isinstance(sample[field], str):
                    if regexp.search(sample[field]):
                        return False
                else:
                    if regexp.search(str(sample[field])):
                        return False
        return True

    if inplace:
        self.dataset = self.dataset.filter(filter_)
        logger.info(
            f"Filtered dataset in-place with regex matching on fields: {fields}"
        )
        # return self for chaining
        return self
    else:
        new_dataset = self.dataset.filter(filter_)
        logger.info(
            f"Filtered dataset with regex matching on fields: {fields}"
        )
        return type(self)(
            new_dataset,
            model=self.model,
            emb_matrix=self.emb_matrix,
            cluster_ids=self.cluster_ids,
            cluster_centers=self.cluster_centers,
            openai_api_key=self.openai_api_key,
        )

src/galactic/filters.py
taylorai/galactic
import logging
import datasets
import pandas as pd
import numpy as np
from typing import Optional, Sequence
from dataclasses import dataclass
from . import (
    classifiers,
    conversations,
    loaders,
    filters,
    taggers,
    transforms,
    embedding,
    clustering,
    # minhash_lsh,
    visualize,
    # scraping,
)
from .base import GalacticDatasetBase

# set up logging
from .logger import setup_logger

setup_logger()
logger = logging.getLogger("galactic")


@dataclass
class GalacticDataset(GalacticDatasetBase):
    """
    This is the GalacticDataset class. It contains methods for loading data,
    filtering data, tagging data, transforming data, and classifying data.
    """

    dataset: datasets.Dataset
    model: Optional[object] = None
    emb_matrix: Optional[np.ndarray] = None
    cluster_ids: Optional[Sequence[int]] = None
    cluster_centers: Optional[dict[int, np.ndarray]] = None
    openai_api_key: Optional[str] = None
    max_tokens_per_minute: Optional[int] = 100000
    max_requests_per_minute: Optional[int] = 2000

    # attach all the imported methods
    ## loaders
    from_csv = classmethod(loaders.from_csv)
    from_jsonl = classmethod(loaders.from_jsonl)
    from_pandas = classmethod(loaders.from_pandas)
    from_parquet = classmethod(loaders.from_parquet)
    from_hugging_face = classmethod(loaders.from_hugging_face)
    from_hugging_face_stream = classmethod(loaders.from_hugging_face_stream)
    save = loaders.save
    ## conversation utils
    conversation_from_dicts = conversations.conversation_from_dicts
    conversation_from_string = conversations.conversation_from_string
    convert_conversation_to_string = (
        conversations.convert_conversation_to_string
    )
    standardize_last_turn = conversations.standardize_last_turn
    get_conversation_length = conversations.get_conversation_length
    get_conversation_speakers = conversations.get_conversation_speakers
    is_alternating = conversations.is_alternating
    get_last_speaker = conversations.get_last_speaker
    get_shared_prefix = conversations.get_shared_prefix
    add_initial_system_message = conversations.add_initial_system_message
    take_initial_system_message = conversations.take_initial_system_message
    take_last_message = conversations.take_last_message
    ## filters
    filter_string = filters.filter_string
    filter_regex = filters.filter_regex
    apply_bloom_filter = filters.apply_bloom_filter
    ## taggers
    tag_string = taggers.tag_string
    tag_regex = taggers.tag_regex
    detect_language = taggers.detect_language
    detect_pii = taggers.detect_pii
    count_tokens = taggers.count_tokens
    calc_perplexity = taggers.calc_perplexity
    ai_tagger = taggers.ai_tagger
    ## transforms
    trim_whitespace = transforms.trim_whitespace
    unicode_normalize = transforms.unicode_normalize
    ai_column = transforms.ai_column
    ## classifiers
    ai_classifier = classifiers.ai_classifier
    fasttext_classifier = classifiers.fasttext_classifier
    embeddings_classifier = classifiers.embeddings_classifier
    train_fasttext_classifier = classifiers.train_fasttext_classifier
    train_embeddings_classifier = classifiers.train_embeddings_classifier
    ## embedding
    initialize_embedding_model = embedding.initialize_embedding_model
    get_embeddings = embedding.get_embeddings
    get_nearest_neighbors = embedding.get_nearest_neighbors
    reduce_embedding_dim = embedding.reduce_embedding_dim
    ## clustering
    cluster = clustering.cluster
    recompute_cluster_centers = clustering.recompute_cluster_centers
    _get_clusters = clustering._get_clusters
    remove_cluster = clustering.remove_cluster
    ai_label_clusters = clustering.ai_label_clusters
    get_cluster_info = clustering.get_cluster_info
    semdedup = clustering.semdedup
    ## minhash lsh

    ## visualizations
    plot_embeddings = visualize.plot_embeddings

    ## scraping
    # scrape_urls = scraping.scrape_urls
    # postprocess_scraped_pages = scraping.postprocess_scraped_pages

    def __post_init__(self):
        """
        Initializes the GalacticDataset instance.
        If '__id' does not exist in dataset columns, it is added.
        If the dataset contains cluster and embedding info but no cluster centers, they are set.
        If the dataset contains embeddings, the embedding matrix is set.
        Raises:
        ValueError: If the dataset contains duplicate __id values.
        """
        # add unique increaing int __id field if it doesn't already exist
        if "__id" not in self.dataset.column_names:
            self.dataset = self.dataset.map(
                lambda _, i: {"__id": i},
                with_indices=True,
            )
        elif "__id" in self.dataset.column_names:
            if len(self.dataset["__id"]) != len(set(self.dataset["__id"])):
                raise ValueError("Dataset contains duplicate __id values.")

    def __repr__(self):
        return pd.DataFrame(self.dataset.select(range(10))).__repr__()

    def __str__(self):
        return self.__repr__()

    def __len__(self):
        return len(self.dataset)

    def __getitem__(self, idx):
        return self.dataset[idx]

    # delegate basic stuff to the underlying Dataset object
    def __getattr__(self, name):
        if name in ["column_names", "features", "info"]:
            return getattr(self.dataset, name)
        elif name in ["filter", "map", "select", "shuffle", "select_columns"]:
            if name == "shuffle":
                logger.warning(
                    "Shuffling the dataset can be really expensive! Watch out!"
                )

            def wrapper(*args, **kwargs):
                result = getattr(self.dataset, name)(*args, **kwargs)
                if name == "shuffle":
                    result = result.flatten_indices()
                return GalacticDataset(
                    dataset=result,
                    model=self.model,
                    emb_matrix=self.emb_matrix,
                    cluster_ids=self.cluster_ids,
                    cluster_centers=self.cluster_centers,
                    openai_api_key=self.openai_api_key,
                )

            return wrapper

    def drop_column(
        self,
        column: str,
    ):
        """
        This method drops a column from the dataset.

        :param column: The column to drop.
        :type column: str
        """
        self.dataset = self.dataset.select_columns(
            [x for x in self.dataset.column_names if x != column]
        )

    def set_openai_key(self, key):
        """
        This method sets the OpenAI API key.

        :param key: The OpenAI API key.
        :type key: str
        """
        self.openai_api_key = key

    def set_rate_limits(self, **kwargs):
        if "max_tokens_per_minute" in kwargs:
            self.max_tokens_per_minute = kwargs["max_tokens_per_minute"]
        if "max_requests_per_minute" in kwargs:
            self.max_requests_per_minute = kwargs["max_requests_per_minute"]

src/galactic/galactic.py
taylorai/galactic
# Based on code from https://huggingface.co/edugp/kenlm/, MIT License

import os
import re
import unicodedata
from typing import Dict

import kenlm
import sentencepiece
from huggingface_hub import hf_hub_download


class SentencePiece:
    def __init__(
        self,
        model: str,
    ):
        super().__init__()
        self.sp = sentencepiece.SentencePieceProcessor()
        self.sp.load(str(model))

    def do(self, text: dict) -> dict:
        tokenized = self.sp.encode_as_pieces(text)
        return " ".join(tokenized)


class KenlmModel:
    digit_re: re.Pattern = re.compile(r"\d")
    unicode_punct: Dict[str, str] = {
        "，": ",",
        "。": ".",
        "、": ",",
        "„": '"',
        "”": '"',
        "“": '"',
        "«": '"',
        "»": '"',
        "１": '"',
        "」": '"',
        "「": '"',
        "《": '"',
        "》": '"',
        "´": "'",
        "∶": ":",
        "：": ":",
        "？": "?",
        "！": "!",
        "（": "(",
        "）": ")",
        "；": ";",
        "–": "-",
        "—": " - ",
        "．": ". ",
        "～": "~",
        "’": "'",
        "…": "...",
        "━": "-",
        "〈": "<",
        "〉": ">",
        "【": "[",
        "】": "]",
        "％": "%",
        "►": "-",
    }
    unicode_punct_re = re.compile(f"[{''.join(unicode_punct.keys())}]")
    non_printing_chars_re = re.compile(
        f"[{''.join(map(chr, list(range(0,32)) + list(range(127,160))))}]"
    )
    kenlm_model_dir = None
    sentence_piece_model_dir = None

    def __init__(
        self,
        model_dataset: str,
        language: str,
        lower_case: bool = False,
        remove_accents: bool = False,
        normalize_numbers: bool = True,
        punctuation: int = 1,
    ):
        # change here: rather than expecting the model to be downloaded, we'll just download it on demand.
        # this is nice because the repo of all the KenLM models is fucking enormous
        hf_repo_id = "edugp/kenlm"

        # model
        model_hf_path = os.path.join(model_dataset, f"{language}.arpa.bin")
        model_local_path = hf_hub_download(hf_repo_id, filename=model_hf_path)
        self.model = kenlm.Model(model_local_path)

        # tokenizer
        tokenizer_hf_path = os.path.join(model_dataset, f"{language}.sp.model")
        tokenizer_local_path = hf_hub_download(
            hf_repo_id, filename=tokenizer_hf_path
        )
        self.tokenizer = SentencePiece(tokenizer_local_path)
        self.accent = remove_accents
        self.case = lower_case
        self.numbers = normalize_numbers
        self.punct = punctuation

    @classmethod
    def from_pretrained(
        cls,
        model_dataset: str,
        language: str,
    ):
        return cls(
            model_dataset,
            language,
            False,
            False,
            True,
            1,
        )

    def pp(self, log_score, length):
        return 10.0 ** (-log_score / length)

    def get_perplexity(self, doc: str, normalize_cc_net: bool = True):
        if normalize_cc_net:
            doc = self.normalize(
                doc,
                accent=self.accent,
                case=self.case,
                numbers=self.numbers,
                punct=self.punct,
            )
        # Tokenize (after normalizing): See https://github.com/facebookresearch/cc_net/blob/bda555bd1cf1ee2e0b925363e62a61cd46c8b60d/cc_net/mine.py#L352 for full pipeline
        doc = self.tokenizer.do(doc)
        doc_log_score, doc_length = 0, 0
        for line in doc.split("\n"):
            log_score = self.model.score(line)
            length = len(line.split()) + 1
            doc_log_score += log_score
            doc_length += length
        return round(self.pp(doc_log_score, doc_length), 1)

    def normalize(
        self,
        line: str,
        accent: bool = True,
        case: bool = True,
        numbers: bool = True,
        punct: int = 1,
    ) -> str:
        line = line.strip()
        if not line:
            return line
        if case:
            line = line.lower()
        if accent:
            line = self.strip_accents(line)
        if numbers:
            line = self.digit_re.sub("0", line)
        if punct == 1:
            line = self.replace_unicode_punct(line)
        elif punct == 2:
            line = self.remove_unicode_punct(line)
        line = self.remove_non_printing_char(line)
        return line

    def strip_accents(self, line: str) -> str:
        """Strips accents from a piece of text."""
        nfd = unicodedata.normalize("NFD", line)
        output = [c for c in nfd if unicodedata.category(c) != "Mn"]
        if len(output) == line:
            return line
        return "".join(output)

    def replace_unicode_punct(self, text: str) -> str:
        return "".join(self.unicode_punct.get(c, c) for c in text)

    def remove_unicode_punct(self, text: str) -> str:
        """More aggressive version of replace_unicode_punct but also faster."""
        return self.unicode_punct_re.sub("", text)

    def remove_non_printing_char(self, text: str) -> str:
        return self.non_printing_chars_re.sub("", text)

src/galactic/kenlm.py
taylorai/galactic
import os
from typing import Optional, Callable
import datasets
import json
import pybloom_live
import pandas as pd
from .utils import read_csv, read_jsonl
from tqdm.auto import tqdm
import pyarrow.parquet as pq

import logging

logger = logging.getLogger("galactic")


def from_csv(cls, path: str):
    """
    Loads a dataset from a CSV file.

    .. code-block:: python

        # Example usage:
        ds = GalacticDataset.from_csv('path_to_file.csv')

    :param path: The path to the CSV file.
    :type path: str
    :return: An instance of the class populated with the dataset from the CSV file.
    :rtype: cls
    """
    df = read_csv(path)
    return cls.from_pandas(df)


def from_jsonl(cls, path):
    """
    Loads a dataset from a JSONL file.

    .. code-block:: python

        # Example usage:
        ds = GalacticDataset.from_jsonl('path_to_file.jsonl')

    :param path: The path to the JSONL file.
    :type path: str
    :return: An instance of the class populated with the dataset from the JSONL file.
    :rtype: cls
    """
    df = read_jsonl(path)
    return cls.from_pandas(df)


def from_parquet(cls, path):
    """
    Loads a dataset from a Parquet file.

    .. code-block:: python

        # Example usage:
        ds = GalacticDataset.from_parquet('path_to_file.parquet')

    :param path: The path to the Parquet file.
    :type path: str
    :return: An instance of the class populated with the dataset from the Parquet file.
    :rtype: cls
    """
    df = pd.read_parquet(path)
    return cls.from_pandas(df)


def from_pandas(cls, df):
    """
    Loads a dataset from a Pandas DataFrame.

    .. code-block:: python

        # Example usage:
        ds = GalacticDataset.from_pandas(dataframe)

    :param df: The Pandas DataFrame containing the dataset.
    :type df: pd.DataFrame
    :return: An instance of the class populated with the dataset from the DataFrame.
    :rtype: cls
    """
    dataset = datasets.Dataset.from_pandas(df)
    return cls(dataset)


def from_hugging_face(
    cls, path: str, split: str, config_name: Optional[str] = None, **kwargs
):
    """
    Loads a dataset from the Hugging Face hub.

    .. code-block:: python

        # Example usage:
        ds = GalacticDataset.from_hugging_face('dataset_path', 'split', 'config_name')

    :param path: The path to the dataset on the Hugging Face hub.
    :type path: str
    :param split: The specific split of the dataset to load.
    :type split: str
    :param config_name: The name of the configuration to load, if applicable. Defaults to None.
    :type config_name: Optional[str]
    :param kwargs: Additional keyword arguments passed to datasets.load_dataset().
    :return: An instance of the class populated with the dataset from the Hugging Face hub.
    :rtype: cls
    """
    dataset = datasets.load_dataset(
        path, name=config_name, split=split, **kwargs
    )
    return cls(dataset)


def from_hugging_face_stream(
    cls,
    path: str,
    split: str,
    config_name: Optional[str] = None,
    filters: list[Callable[[dict], bool]] = None,
    dedup_fields: Optional[list[str]] = None,
    max_samples: Optional[int] = 200000,
    **kwargs,
):
    """
    Loads a dataset from the Hugging Face hub, streaming the data from disk.

    .. code-block:: python

        # Example usage:
        ds = GalacticDataset.from_hugging_face_stream('dataset_path', 'split', filters=[filter_func], dedup_fields=['field1', 'field2'])

    :param path: The path to the dataset on the Hugging Face hub.
    :type path: str
    :param split: The specific split of the dataset to load.
    :type split: str
    :param config_name: The name of the configuration to load, if applicable. Defaults to None.
    :type config_name: Optional[str]
    :param filters: A list of filter functions to apply to the stream. Each filter function should take a dictionary and return a boolean. Defaults to None.
    :type filters: Optional[list[Callable[[dict], bool]]]
    :param dedup_fields: A list of fields to use for deduplication via a Bloom filter. Defaults to None.
    :type dedup_fields: Optional[list[str]]
    :param max_samples: The maximum number of samples to load. Defaults to 200000.
    :type max_samples: Optional[int]
    :param kwargs: Additional keyword arguments passed to datasets.load_dataset().
    :return: An instance of the class populated with the streamed dataset from the Hugging Face hub.
    :rtype: cls
    """
    handle = datasets.load_dataset(
        path, name=config_name, split=split, streaming=True, **kwargs
    )
    # first, apply the filters, which will operate on the stream
    if filters is not None:
        for filter in filters:
            handle = handle.filter(filter)

    # then set up bloom filter if applicable
    bloom = None
    if dedup_fields is not None and len(dedup_fields) > 0:
        bloom = pybloom_live.BloomFilter(
            capacity=max_samples if max_samples is not None else int(1.0e9),
            error_rate=0.001,
        )

    samples = []
    progress = tqdm(total=max_samples, desc=f"Loading samples from {path}")
    for idx, sample in enumerate(handle):
        if bloom is not None:
            key = json.dumps({k: sample[k] for k in dedup_fields})
            if key in bloom:
                continue
            bloom.add(key)
        samples.append(sample)
        progress.update(1)
        if len(samples) >= max_samples:
            break
    progress.close()
    msg = f"Loaded {len(samples)} samples from {path}. "
    if idx > len(samples):
        msg += f"Removed {idx - len(samples)} samples that were duplicated. "
    logger.info(msg)
    return cls(datasets.Dataset.from_list(samples))


# save dataset as jsonl or csv
def save(self, path: str, overwrite: bool = False) -> None:
    """
    Saves the dataset as a CSV or JSONL file.

    .. code-block:: python

        # Example usage:
        ds.save('path_to_file.csv', overwrite=True)

    :param path: The path where the dataset will be saved. The format is inferred from the file extension (either .csv or .jsonl).
    :type path: str
    :param overwrite: Whether to overwrite the file if it already exists. Defaults to False.
    :type overwrite: bool
    :raises ValueError: If the path already exists and overwrite is False, or if the file format is unsupported.
    """
    # check if exists
    if os.path.exists(path) and not overwrite:
        raise ValueError(
            f"Path {path} already exists. Use overwrite=True to overwrite."
        )
    # get format
    ext = path.split(".")[-1]
    if ext not in ["csv", "jsonl"]:
        raise ValueError(
            f"Unsupported file format {ext}. Must be either csv or jsonl."
        )

    # save
    if ext == "csv":
        self.dataset.to_csv(path)
    elif ext == "jsonl":
        df = self.dataset.to_pandas()
        df.to_json(path, orient="records", lines=True)

src/galactic/loaders.py
taylorai/galactic
import logging
from colorama import init, Fore, Style

init(autoreset=True)
import logging
from colorama import init, Fore, Style


def setup_logger():
    init(autoreset=True)

    class ColoredFormatter(logging.Formatter):
        def format(self, record):
            colors = {
                logging.DEBUG: Fore.CYAN,
                logging.INFO: Fore.GREEN,
                logging.WARNING: Fore.YELLOW,
                logging.ERROR: Fore.RED,
                logging.CRITICAL: Fore.RED,
            }
            levelname_color = colors.get(record.levelno, Fore.WHITE)
            record.levelname = (
                f"{levelname_color}{record.levelname}{Style.RESET_ALL}"
            )
            return super().format(record)

    # Initialize logger
    logger = logging.getLogger("galactic")
    logger.setLevel(logging.INFO)

    # Initialize and set handler and formatter
    colored_formatter = ColoredFormatter("%(levelname)s: %(message)s")
    console_handler = logging.StreamHandler()
    console_handler.setFormatter(colored_formatter)
    logger.addHandler(console_handler)

    # Suppress third-party logs
    third_party_loggers = [
        "urllib3",
        "filelock",
        "git",
        "tensorflow",
        "jax",
        "h5py",
        "jaxlib",
        "hpack",
        "asyncio",
        "modal-utils",
        "fsspec",
        "httpx",
    ]
    for log_name in third_party_loggers:
        logging.getLogger(log_name).setLevel(logging.WARNING)

    # do not propagate
    logger.propagate = False

src/galactic/logger.py
taylorai/galactic
import re
import unicodedata
from datasketch import MinHash, LeanMinHash

from logging import getLogger

logger = getLogger("galactic")


def replace_whitespace_with_underscore(text: str):
    return re.sub(r"\s+", "_", text)


def compute_minhash(text: str, k: int = 9, num_perm: int = 128) -> LeanMinHash:
    # first normalize
    text = unicodedata.normalize("NFKC", text)

    # then replace whitespace with underscore
    text = replace_whitespace_with_underscore(text)

    # shingle
    shingles = set()
    for i in range(len(text) - k + 1):
        shingles.add(text[i : i + k])

    # compute minhash
    m = MinHash(num_perm=num_perm)
    for s in shingles:
        m.update(s.encode("utf8"))
    return LeanMinHash(m)


def compute_minhashes(self, field: str, k: int = 9, num_perm: int = 128):
    # make sure field exists
    if field not in self.dataset.features:
        raise ValueError(f"Field {field} not found in dataset.")

    def minhash_(sample):
        m = compute_minhash(sample[field], k=k, num_perm=num_perm)
        buf = bytearray(m.bytesize())
        m.serialize(buf)
        return {f"__minhash__{field}": buf}

    self.dataset = self.dataset.map(minhash_)
    logger.info(
        f"Computed minhashes for field {field}, added minhash metadata to '__minhash__{field}'."
    )
    return self

src/galactic/minhash_lsh.py
taylorai/galactic
# basically what we're gonna do here is give the user a tool to go from column of URLs to column of scraped pages,
# using tool of choice (requests, selenium, puppeteer, trafilatura, etc.)
from typing import Literal
import requests
import justext
import trafilatura
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from webdriver_manager.chrome import ChromeDriverManager
from readabilipy import simple_json_from_html_string
from boilerpy3 import extractors

extractor = extractors.ArticleExtractor()


def scrape_urls(
    self,
    url_column: str,
    output_column: str,
    backend: Literal["requests", "trafilatura"],
):
    if backend == "selenium":
        # Initialize Selenium in headless mode
        chrome_options = webdriver.ChromeOptions()
        chrome_options.add_argument("--headless")
        driver_service = Service(ChromeDriverManager().install())
        driver = webdriver.Chrome(
            service=driver_service, options=chrome_options
        )

        def _scrape_single_url(url):
            driver.get(url)
            return driver.page_source

    elif backend == "requests":

        def _scrape_single_url(url):
            try:
                response = requests.get(url, timeout=5)
                if response.status_code == 200:
                    return response.text
                else:
                    return "Failed to scrape"
            except Exception as e:
                return f"Failed to scrape."

    elif backend == "trafilatura":

        def _scrape_single_url(url):
            downloaded = trafilatura.fetch_url(url)
            result = trafilatura.extract(downloaded)
            if result:
                return result
            else:
                return "Failed to scrape"

    else:
        raise ValueError("Invalid backend")

    self.dataset = self.dataset.map(
        lambda sample: {output_column: _scrape_single_url(sample[url_column])}
    )


def postprocess_scraped_pages(
    self,
    input_column: str,
    output_column: str,
    backend: Literal[
        "readabilipy", "justext", "boilerpy3", "goose3"
    ] = "readabilipy",
):
    if backend == "readabilipy":

        def _postprocess_single_page(page):
            json_page = simple_json_from_html_string(page)
            result = ""
            if json_page["title"]:
                result += f"Title: {json_page['title']}\n\n"
            if json_page["byline"]:
                result += f"By: {json_page['byline']}\n\n"
            if json_page["plain_content"]:
                result += json_page["plain_content"]

            return result

    elif backend == "justext":

        def _postprocess_single_page(page):
            ## TODO: detect the language automatically
            paragraphs = justext.justext(page, justext.get_stoplist("English"))
            result = ""
            for paragraph in paragraphs:
                if not paragraph.is_boilerplate:
                    result += paragraph.text + "\n\n"

            return result

    elif backend == "boilerpy3":

        def _postprocess_single_page(page):
            if len(page) == 0:
                return "Failed to scrape"
            try:
                doc = extractor.get_doc(page)
                content = doc.content
                title = doc.title
                return f"Title: {title}\n\n{content}"
            except Exception as e:
                return "Failed to scrape"

    elif backend == "goose3":
        from goose3 import Goose

        g = Goose()

        def _postprocess_single_page(page):
            try:
                article = g.extract(raw_html=page)
                return article.cleaned_text
            except Exception as e:
                return "Failed to scrape"

    else:
        raise ValueError("Invalid postprocess")

    self.dataset = self.dataset.map(
        lambda sample: {
            output_column: _postprocess_single_page(sample[input_column])
        }
    )

    return self

src/galactic/scraping.py
taylorai/galactic
import os
import re
import pathlib
from typing import Sequence, Optional, Union, Literal
import numpy as np
import huggingface_hub
from transformers import AutoTokenizer
import scrubadub
import tiktoken
import jinja2
import time

from .utils import byte_len
from .async_openai import run_chat_queries_with_openai

import logging

logger = logging.getLogger("galactic")


def tag_string(self, fields: Sequence[str], values: Sequence[str], tag: str):
    """Tag data containing a particular string or list of strings in the specified field."""
    # make sure the tag hasn't already been used
    if f"__tag__{tag}" in self.dataset.column_names:
        logger.warning(
            f"Tag {tag} already exists. This will overwrite the existing tag."
        )

    regexp = re.compile("|".join([re.escape(val) for val in values]))

    def tag_(sample):
        for field in fields:
            if field in sample:
                if isinstance(sample[field], str):
                    if regexp.search(sample[field]):
                        return {f"__tag__{tag}": True}
                else:
                    if regexp.search(str(sample[field])):
                        return {f"__tag__{tag}": True}
        return {f"__tag__{tag}": False}

    self.dataset = self.dataset.map(tag_)
    logger.info(
        f"Tagged dataset in-place with exact string matching on fields: {fields}"
    )
    # return self for chaining
    return self


def tag_regex(self, fields: Sequence[str], regex: str, tag: str):
    # make sure the tag hasn't already been used
    if f"__tag__{tag}" in self.dataset.column_names:
        logger.warning(
            f"Tag already exists. This will overwrite the existing tag."
        )

    regexp = re.compile(regex)

    def tag_(sample):
        for field in fields:
            if field in sample:
                if isinstance(sample[field], str):
                    if regexp.search(sample[field]):
                        return {f"__tag__{tag}": True}
                else:
                    if regexp.search(str(sample[field])):
                        return {f"__tag__{tag}": True}
        return {f"__tag__{tag}": False}

    self.dataset = self.dataset.map(tag_)
    logger.info(
        f"Tagged dataset in-place with tag '__tag__{tag}', using regex matching on fields: {fields}"
    )
    return self


def detect_language(
    self, field: str, method: Literal["fasttext", "langdetect"] = "fasttext"
):
    """Detect the language of the specified field."""
    # make sure field exists
    if field not in self.dataset.features:
        raise ValueError(f"Field {field} not found in dataset.")

    if method == "fasttext":
        import fasttext

        model_path = huggingface_hub.hf_hub_download(
            repo_id="TaylorAI/galactic-models", filename="lid.176.ftz"
        )
        model = fasttext.load_model(model_path)

        def detect_(sample):
            if isinstance(sample[field], str):
                return {
                    "__language": model.predict(
                        sample[field].replace("\n", " ")
                    )[0][0].split("__label__")[1]
                }
            else:
                return {
                    "__language": model.predict(str(sample[field]))[0][
                        0
                    ].split("__label__")[1]
                }

    elif method == "langdetect":
        import langdetect

        def detect_(sample):
            if isinstance(sample[field], str):
                return {"__language": langdetect.detect(sample[field])}
            else:
                return {"__language": langdetect.detect(str(sample[field]))}

    self.dataset = self.dataset.map(detect_)
    logger.info(
        f"Detected language in field {field}, added language metadata to '__language'."
    )
    return self


def calc_perplexity(
    self,
    field: str,
    model: str = "kenlm",  # other option is pythia
    language: Optional[str] = "en",
    dataset: Optional[str] = "wikipedia",
):
    """
    Calculate the perplexity-per-byte of the specified field.

    .. note::

        You must install KenLM to calculate perplexity using the KenLM model from https://github.com/kpu/kenlm.
        To install, run ``pip install https://github.com/kpu/kenlm/archive/master.zip``.

    .. code-block:: python

        # Replace "field_name1"
        ds.calc_perplexity(field="field_name1")

    """
    # make sure field exists and is a string field
    if field not in self.dataset.features:
        raise ValueError(f"Field {field} not found in dataset.")
    elif self.dataset.features[field].dtype != "string":
        raise ValueError(
            f"Field {field} is not a string field, and so can't be used to calculate perplexity."
        )
    if model == "pythia":
        import ctranslate2

        repo_path = huggingface_hub.snapshot_download(
            "TaylorAI/galactic-models", allow_patterns="p70/*"
        )
        model_path = pathlib.Path(repo_path) / "p70"
        model = ctranslate2.Generator(str(model_path))
        tokenizer = AutoTokenizer.from_pretrained("EleutherAI/pythia-70m")

        def calc_(sample):
            token_ids = tokenizer(sample[field]).input_ids
            tokens = tokenizer.convert_ids_to_tokens(token_ids)
            log_probs = model.score_batch([tokens])[0].log_probs
            ppl = np.exp(-np.sum(log_probs) / byte_len(sample[field]))
            return {"__perplexity": ppl}

    elif model == "kenlm":
        try:
            from .kenlm import KenlmModel
        except ImportError:
            raise ImportError(
                "KenLM is not installed. Install with 'pip install https://github.com/kpu/kenlm/archive/master.zip'."
            )
        if language is None or dataset is None:
            raise ValueError(
                "Must specify language (e.g. 'en') and dataset (e.g. 'wikipedia') for KenLM. See options here: https://huggingface.co/edugp/kenlm/tree/main"
            )
        model = KenlmModel.from_pretrained(dataset, language)

        def calc_(sample):
            ppl = model.get_perplexity(sample[field])
            return {"__perplexity": ppl}

    else:
        raise ValueError(
            f"Model {model} not supported. Supported models: 'kenlm', 'pythia'."
        )

    self.dataset = self.dataset.map(calc_)
    logger.info(
        f"Calculated perplexity-per-byte in field {field}, added perplexity metadata to '__perplexity'."
    )
    return self


def detect_pii(self, fields: Sequence[str]):
    """
    Detect personally identifiable information in the specified fields.
    Args:
        fields (List[str]): List of fields to detect PII in.
        Currently only supports "email", "phone", and "credential".
    """
    # make sure all the fields exist
    for field in fields:
        if field not in self.dataset.features:
            raise ValueError(f"Field {field} not found in dataset.")

    def detect_(sample):
        text = " ".join([str(sample[field]) for field in fields])
        filth = scrubadub.list_filth(text)
        filth_types = [f.detector_name for f in filth]
        return {
            **{
                f"__pii__{f}": f in filth_types
                for f in ["email", "phone", "credential"]
            },
            "__pii__any": len(filth) > 0,
        }

    self.dataset = self.dataset.map(detect_)
    logger.info(
        f"Detected PII in fields: {fields}; added __pii__email, __pii__phone, __pii__credential, and __pii__any metadata."
    )
    # no option to do out-of-place as this operation is not destructive
    return self


def count_tokens(self, fields: Sequence[str], tokenizer: Optional[str] = None):
    """
    Count the number of tokens in the specified fields.

    .. code-block:: python

        # Replace "field_name1" and "field_name2" with your fields
        ds.count_tokens(fields=["field_name1, field_name2"])

    Args:
        fields (List[str]): List of fields to count tokens in.
        tokenizer (Callable): Tokenizer function to use. Defaults to None, which uses bytes.
    """
    if tokenizer is None:
        # count bytes in string
        count_fn = lambda x: byte_len(str(x))
        field_name = "__byte_count__"
    else:
        tokenizer = AutoTokenizer.from_pretrained(tokenizer)
        count_fn = lambda x: len(tokenizer(str(x)).input_ids)
        field_name = "__token_count__"
    self.dataset = self.dataset.map(
        lambda x: {
            f"{field_name}{field}": count_fn(x[field]) for field in fields
        }
    )
    logger.info(
        f"Counted tokens in fields: {fields}, added metadata to {field_name}"
    )

    # no option to do out-of-place as this operation is not destructive
    return self


def ai_tagger(
    self,
    field: str,
    tags: Union[list[str], dict[str, str]],
    prompt: Optional[str] = None,
    backend="openai",
    allow_not_sure: bool = False,
):
    """
    Use OpenAI's API or a HF zero-shot classifier to generate a new column based on an existing column.
    Args:
        name (str): Name of the new column.
        field (str): Name of the field to use for classification. Can be None if backend is 'embeddings'.
        classes (List[str] or dict[str, str]): Classes to classify into. Can be just a list of labels, or a dict of label: description.
        If provided, descriptions will be used for a) prompting API model if backend = 'openai', b) embedding if backend = "embeddings",
        c) zero-shot classification if backend = "huggingface".
        backend (str): Which backend to use. Currently supports 'embeddings', 'openai' and 'huggingface'.
    """
    if backend != "openai":
        raise NotImplementedError(
            "Only OpenAI API is currently supported for AI tagging."
        )

    elif backend == "openai":
        if self.openai_api_key is None:
            self.openai_api_key = os.environ.get("OPENAI_API_KEY", None)
            if self.openai_api_key is None:
                raise ValueError(
                    "No OpenAI API key found. Set OPENAI_API_KEY environment variable, or set the openai_api_key attribute of the GalacticDataset."
                )
        # make sure field exists
        if field not in self.dataset.column_names:
            raise ValueError(f"Field {field} not found in dataset.")

        # create logit bias (same for each tag) -- model can only output True or False
        token2cls = {}
        tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")
        classes = ["True", "False"]
        if allow_not_sure:
            classes.append("Not sure")
        for cls in classes:
            t1 = tokenizer.encode(cls)[0]
            t2 = tokenizer.encode(" " + cls)[0]
            if t1 in token2cls or t2 in token2cls:
                raise ValueError(
                    "Class names must not share a prefix for logit bias trick to work. Try changing the labels and see if that helps."
                )
            else:
                token2cls[t1] = cls
                token2cls[t2] = cls
        logit_bias = {t: 100 for t in token2cls.keys()}

        for idx, tag in enumerate(tags):
            logger.info(f"Tagging with tag {tag}...")
            # construct prompt template
            if prompt is None:
                prompt_template = (
                    f"Tag the provided text with the following tag:\n"
                )
                if isinstance(tags, list):
                    prompt_template += f"  - {tag}\n"
                else:
                    prompt_template += f"  - {tag}: {tags[tag]}\n"
                prompt_template += "Answer True if the tag applies to the text, and False if it does not. "
                if allow_not_sure:
                    prompt_template += "If you're not sure, answer Not sure."
                prompt_template += (
                    "\n\n---\n\nText: {{"
                    + field
                    + "}}\n\n---\n\nDoes the tag apply?"
                )
            else:
                prompt_template = prompt
            prompt_template = jinja2.Template(prompt_template)
            prompts = self.dataset.select_columns([field]).map(
                lambda sample: {"__prompt": prompt_template.render(**sample)}
            )["__prompt"]
            logger.info(f"Example prompt for tag {tag}: {prompts[0]}")

            responses = run_chat_queries_with_openai(
                queries=prompts,
                api_key=self.openai_api_key,
                logit_bias=logit_bias,
                max_new_tokens=1,
                max_requests_per_minute=self.max_requests_per_minute,
                max_tokens_per_minute=self.max_tokens_per_minute,
            )

            # map back to labels
            first_tokens = [
                tokenizer.encode(response) for response in responses
            ]
            labels = [token2cls[t[0]] for t in first_tokens]
            self.dataset = self.dataset.add_column(
                name=f"__ai_tag__{tag.replace(' ', '_')}", column=labels
            )
            logger.info(f"Tagged with tag {tag}.")
            # if not the last tag, cool down for 30 seconds
            if idx < len(tags) - 1:
                logger.info(
                    f"Cooling down for 30 seconds before tagging with next tag."
                )
                time.sleep(30)

    return self

src/galactic/taggers.py
taylorai/galactic
# a transform can either modify a column in-place, or add a new column. if the new column is just true/false, it should probably be a "tag".
import os
import unicodedata
from typing import Sequence
import logging
import jinja2
from typing import Optional
from .async_openai import run_chat_queries_with_openai

logger = logging.getLogger("galactic")


def trim_whitespace(self, fields: Sequence[str], inplace: bool = True):
    """
    Trim Unicode-defined whitespace at the beginning and end of specified fields.

    :param fields: List of fields to trim.
    :type fields: Sequence[str]
    :param inplace: Whether to perform the operation in-place, defaults to True
    :type inplace: bool, optional
    :return: Modified GalacticDataset instance.
    :rtype: GalacticDataset

    .. code-block:: python

        dataset = GalacticDataset(…)
        dataset.trim_whitespace(fields=['column1', 'column2'])

    """

    def trim_(sample):
        for field in fields:
            if field in sample:
                sample[field] = sample[field].strip()
        return sample

    if inplace:
        self.dataset = self.dataset.map(trim_)
        logger.info(f"Trimmed whitespace for fields: {fields}")
        # return self for chaining
        return self
    else:
        new_dataset = self.dataset.map(trim_)
        logger.info(f"Trimmed whitespace for fields: {fields}")
        return type(self)(
            new_dataset,
            model=self.model,
            emb_matrix=self.emb_matrix,
            cluster_ids=self.cluster_ids,
            cluster_centers=self.cluster_centers,
            openai_api_key=self.openai_api_key,
        )


def unicode_normalize(
    self,
    fields: Sequence[str],
    form: str = "NFC",
    inplace: bool = True,
):
    """
    Apply Unicode normalization to specified fields. This is useful as a pre-processing step for ML training.

    :param fields: List of fields to normalize.
    :type fields: Sequence[str]
    :param form: Unicode normalization form, defaults to "NFC"
    :type form: str, optional
    :param inplace: Whether to perform the operation in-place, defaults to True
    :type inplace: bool, optional
    :return: Modified GalacticDataset instance if inplace is True, else a new instance.
    :rtype: GalacticDataset

    .. code-block:: python

        dataset = GalacticDataset(…)
        dataset.unicode_normalize(fields=['column1', 'column2'], form='NFC')

    """
    # make sure fields exist and are text fields
    for field in fields:
        if field not in self.dataset.features:
            raise ValueError(f"Field {field} not found in dataset.")
        elif self.dataset.features[field].dtype != "string":
            raise ValueError(
                f"Field {field} is not a string field, and so can't be normalized."
            )
    if inplace:
        self.dataset = self.dataset.map(
            lambda sample: {
                field: unicodedata.normalize(form, sample[field])
                for field in fields
            }
        )
        logger.info(f"Normalized fields {fields} using form {form}")
        # return self for chaining
        return self
    else:
        new_dataset = self.dataset.map(
            lambda sample: {
                field: unicodedata.normalize(form, sample[field])
                for field in fields
            }
        )
        logger.info(f"Normalized fields {fields} using form {form}")
        return type(self)(
            new_dataset,
            model=self.model,
            emb_matrix=self.emb_matrix,
            cluster_ids=self.cluster_ids,
            cluster_centers=self.cluster_centers,
            openai_api_key=self.openai_api_key,
        )


def ai_column(
    self,
    new_column: str,
    prompt: str,
    depends_on=list[str],
    normalize: list[str] = ["strip"],  # must be methods of str class
    system_prompt: Optional[str] = None,
):
    """
    Use OpenAI's API to generate a new column based on an existing column.

    :param new_column: Name of the new column.
    :type new_column: str
    :param prompt: Prompt template (Jinja2) to use for the API request.
    :type prompt: str
    :param depends_on: List of fields needed to fill in the template, defaults to an empty list.
    :type depends_on: list[str], optional
    :param normalize: List of string methods to apply for normalization, defaults to ["strip", "lower"]
    :type normalize: list[str], optional
    :param system_prompt: System instructions (optional, same for every request), defaults to None
    :type system_prompt: str, optional
    :return: Modified GalacticDataset instance.
    :rtype: GalacticDataset

    :raises ValueError: If a specified field is not found in the dataset or if the new_column already exists in the dataset.

    .. code-block:: python

        dataset = GalacticDataset(…)
        dataset.ai_column(new_column='ai_column', prompt='{{column1}} is related to {{column2}}', depends_on=['column1', 'column2'])

    """
    for field in depends_on:
        if field not in self.dataset.column_names:
            raise ValueError(f"Field {field} not found in dataset.")
    if new_column in self.dataset.column_names:
        raise ValueError(
            f"Column {new_column} already exists in dataset. Please choose a different name, or drop the column."
        )
    if self.openai_api_key is None:
        self.openai_api_key = os.environ.get("OPENAI_API_KEY", None)
        if self.openai_api_key is None:
            raise ValueError(
                "No OpenAI API key found. Set OPENAI_API_KEY environment variable, or set the openai_api_key attribute of the GalacticDataset."
            )
    template = jinja2.Template(prompt)
    prompts = self.dataset.select_columns(depends_on).map(
        lambda sample: {"__prompt": template.render(**sample)}
    )["__prompt"]
    logger.info(f"Example prompt: {prompts[0]}")
    responses = run_chat_queries_with_openai(
        queries=prompts,
        api_key=self.openai_api_key,
        system_prompt=system_prompt,
        max_tokens_per_minute=self.max_tokens_per_minute,
        max_requests_per_minute=self.max_requests_per_minute,
    )
    # apply normalizations
    responses = [r if r is not None else "" for r in responses]
    for fn in normalize:
        try:
            responses = [getattr(str, fn)(response) for response in responses]
        except AttributeError:
            raise ValueError(f"Unknown normalization function: {fn}")

    # inplace only because this isn't a destructive operation
    self.dataset = self.dataset.add_column(name=new_column, column=responses)
    logger.info(f"Added new column {new_column} using prompt {prompt}")
    # return self for chaining
    return self

src/galactic/transforms.py
taylorai/galactic
import logging
import requests
import pandas as pd
from dataclasses import dataclass

logger = logging.getLogger("galactic")


@dataclass
class MessageConfig:
    """
    This is the MessageConfig class. It contains the configuration for how a chat message is laid out.
    """

    role_field: str
    content_field: str
    user_role: str
    assistant_role: str
    system_role: str


OPENAI_MESSAGE_CONFIG = MessageConfig(
    role_field="role",
    content_field="content",
    user_role="user",
    assistant_role="assistant",
    system_role="system",
)

SHAREGPT_MESSAGE_CONFIG = MessageConfig(
    role_field="from",
    content_field="value",
    user_role="human",
    assistant_role="gpt",
    system_role="system",
)


def byte_len(x: str):
    return len(x.encode("utf-8"))


def handle_redirects(url: str):
    # handle 308 redirects, pandas burps on them
    res = requests.head(url)
    if res.status_code == 308:
        url = res.headers["Location"]
    elif res.status_code != 200:
        raise ValueError(
            f"Could not read dataset from {url}. Status code: {res.status_code}."
        )
    return url


def read_csv(path: str):
    # if it's remote, follow the redirect first
    if path.startswith("http"):
        path = handle_redirects(path)
    # try to read it
    try:
        df = pd.read_csv(path)
        return df
    except Exception as e:
        logger.warning(
            "Unable to parse CSV file. Falling back to latin-1 encoding."
        )
        df = pd.read_csv(path, encoding="latin-1")
        return df


def read_jsonl(path: str):
    # if it's remote, follow the redirect first
    if path.startswith("http"):
        path = handle_redirects(path)
    # try to read it
    df = pd.read_json(path, lines=True, orient="records")
    return df

src/galactic/utils.py
taylorai/galactic
import re
from typing import Optional, Literal
import matplotlib.pyplot as plt
import seaborn as sns

# from matplotlib.colors import ListedColormap
import numpy as np

import logging

logger = logging.getLogger("galactic")


def plot_embeddings(
    self,
    embedding_field: str,
    color_by: Optional[str] = None,
    save_path: Optional[str] = None,
    dot_size: int = 2,
    theme: Literal[
        "darkgrid", "whitegrid", "dark", "white", "ticks"
    ] = "whitegrid",
    width: Optional[int] = None,
    height: Optional[int] = None,
):
    # Check the dimensionality
    if len(self.dataset[embedding_field][0]) != 2:
        raise ValueError("The dimensionality of the embeddings must be 2.")

    # Separate X and Y coordinates
    x_coords, y_coords = zip(*self.dataset[embedding_field])

    # if width and height provided, set them
    if width and height:
        plt.figure(figsize=(width, height))

    sns.set(style=theme)

    if color_by:
        labels = self.dataset[color_by]
        unique_labels = list(set(labels))
        palette = sns.color_palette("husl", len(unique_labels))

        for i, label in enumerate(unique_labels):
            idx = [j for j, x in enumerate(labels) if x == label]
            ax = sns.scatterplot(
                x=np.array(x_coords)[idx],
                y=np.array(y_coords)[idx],
                color=palette[i],
                label=label,
                s=dot_size,
            )

        # Move legend outside the plot
        plt.legend(
            title=color_by,
            fontsize="small",
            markerscale=2,
            bbox_to_anchor=(1, 1),
            loc="upper left",
        )

        plt.tight_layout(rect=[0, 0, 0.85, 1])

    else:
        sns.scatterplot(x=x_coords, y=y_coords, s=dot_size)

    plt.xlabel("Dimension 1")
    plt.ylabel("Dimension 2")
    escaped_color_by = re.sub(r"([#\$%&~_^\\])", r"\\\1", color_by)
    plt.title(
        "2D embeddings, colored by  $\\mathtt{" + escaped_color_by + "}$"
    )

    if save_path:
        plt.savefig(save_path)
        logger.info(f"Saved plot to {save_path}")
    else:
        plt.show()

src/galactic/visualize.py
taylorai/galactic