Upload folder using huggingface_hub
9823a7e verified 9 months ago
28 kB
import abc
import logging
import os
from functools import partial
from pathlib import Path
from typing import Any, Callable, Optional, Sequence, Type, Union

import numpy as np
import pandas as pd

from .. import completion_parsers, constants, processors, utils
from ..decoders import get_fn_completions

CURRENT_DIR = Path(__file__).parent
logging.getLogger().setLevel(logging.INFO)

__all__ = ["BaseAnnotator", "BaseAnnotatorJSON", "SingleAnnotator"]


class BaseAnnotator(abc.ABC):
    """Base class for a pool of annotators.
    Parameters
    ----------
    annotators_config : Path or list of dict, optional
        A dictionary or path to a yaml file containing the configuration for the pool of annotators. If a directory,
        we search for 'configs.yaml' in it. The keys in the first  dictionary should be the annotator's name, and
        the value should be a dictionary of the annotator's configuration which should have the following keys:
        The path is relative to `base_dir` directory.
        - prompt_template (str): a prompt template or path to it. The template should contain placeholders for keys in
            the example dictionary, typically {instruction} and {output_1} {output_2}.
        - fn_completions (str): function in `alpaca_farm.decoders` for completions. Needs to accept as first argument
            `prompts` which is a list of string.
        - completions_kwargs (dict): kwargs for fn_completions. E.g. model_name, max_tokens, temperature,
        tokens_to_avoid
        - fn_completion_parser (str) : Function in `completion_parsers.py` to use for parsing the completions into
        annotations.
        - completion_parser_kwargs (dict) : Kwargs for fn_completion_parser.
        - other kwargs to `SingleAnnotator` such as batch_size
    seed : int, optional
        Seed for the random number generator.
    is_avoid_reannotations : bool, optional
        Whether to avoid re-annotating examples that have already been annotated by the annotator. This will decrease
        cost but can be slightly slower if there are no annotations that can be reused.
    primary_keys : sequence of str, optional
        Keys use to distinguish the example.
    other_keys_to_keep : sequence of str, optional
        Other columns to store besides the annotations.
    is_store_missing_annotations : bool, optional
        Whether to store missing annotations. If True it avoids trying to reannotate examples that have errors.
    base_dir : Path, optional
        Path to the directory containing the annotators configs. I.e. annotators_config will be relative
        to this directory. If None uses self.DEFAULT_BASE_DIR
    is_raise_if_missing_primary_keys : bool, optional
        Whether to ensure that the primary keys are in the example dictionary. If True, raises an error.
    tmp_missing_annototation : Any, optional
        Temporary value to use for missing annotations when `is_store_missing_annotations` is True.
    annotation_type : type, optional
        Type to use for storing the annotations. If None, uses `self.DEFAULT_ANNOTATION_TYPE`.
    """

    DEFAULT_BASE_DIR = constants.EVALUATORS_CONFIG_DIR
    annotator_column = "annotator"
    TMP_MISSING_ANNOTATION = -1
    DEFAULT_ANNOTATION_TYPE = int

    def __init__(
        self,
        primary_keys: Sequence[str],
        annotators_config: Union[utils.AnyPath, list[dict[str, Any]]] = "claude",
        seed: Optional[int] = 0,
        is_avoid_reannotations: bool = True,
        other_keys_to_keep: Sequence[str] = (
            "price_per_example",
            "time_per_example",
            "raw_completion",
        ),
        is_store_missing_annotations: bool = True,
        base_dir: Optional[utils.AnyPath] = None,
        is_raise_if_missing_primary_keys: bool = True,
        annotation_type: Optional[Type] = None,
    ):
        logging.info(f"Creating the annotator from `{annotators_config}`.")
        self.base_dir = Path(base_dir or self.DEFAULT_BASE_DIR)
        self.seed = seed
        self.is_avoid_reannotations = is_avoid_reannotations
        self.primary_keys = list(primary_keys)
        self.all_keys = self.primary_keys + [self.annotator_column]
        self.other_keys_to_keep = list(other_keys_to_keep)
        self.is_store_missing_annotations = is_store_missing_annotations
        self.is_raise_if_missing_primary_keys = is_raise_if_missing_primary_keys
        self.annotation_type = annotation_type or self.DEFAULT_ANNOTATION_TYPE

        self.annotators_config = self._initialize_annotators_config(annotators_config)
        self.annotators = self._initialize_annotators()
        self.df_annotations = None

    ### Abstract methods ###

    @property
    @abc.abstractmethod
    def SingleAnnotator(self) -> Type["SingleAnnotator"]:
        """Class to use for each single annotator."""
        pass

    #######################
    @property
    def available_fields_to_format(self):
        """Fields that can be formatted in the prompt template."""
        return self.all_keys

    @property
    def annotation_key(self) -> str:
        """How to refer to the annotations, this will be the key for annotations in the output."""
        return "annotation"

    @property
    def random_seed_key(self) -> list[str]:
        """What key / column to seed on for the random generator."""
        return list(self.primary_keys)

    ### Public methods ###
    @property
    def annotator_name(self) -> str:
        return Path(self.annotators_config).parent.name

    def __call__(
        self,
        to_annotate: utils.AnyData,
        **decoding_kwargs,
    ) -> list[dict[str, Any]]:
        """Main function for annotating.
        Parameters
        ----------
        to_annotate : list of dict or dataframe
            Examples to annotate. Each dictionary (or row) should contain all of `self.primary_keys`.
        **decoding_kwargs :
            Additional arguments to pass to `fn_completions`.
        Returns
        -------
        annotated : list of dict
            The annotated examples. Each dict will contain all of `self.primary_keys` and `self.annotation_key`.
        """
        if len(to_annotate) == 0:
            return []

        df_to_annotate = self._preprocess(to_annotate)
        df_annotated = self._annotate(df_to_annotate, **decoding_kwargs)
        annotated = self._postprocess_and_store_(df_annotated, to_annotate)
        return annotated

    #######################

    ### Private methods ###
    def _initialize_annotators_config(self, annotators_config):
        # setting it relative to the config directory
        annotators_config = self.base_dir / annotators_config

        if annotators_config.is_dir():
            annotators_config = annotators_config / "configs.yaml"

        return annotators_config

    def _initialize_annotators(self) -> dict[str, "SingleAnnotator"]:
        """Load all the configs and prompts if necessary."""
        annotators_config = utils.load_configs(self.annotators_config)
        try:
            # in case a path is given we make it relative to that path
            base_dir = self.annotators_config.parents[1]
        except:
            base_dir = self.base_dir

        return {
            name: self.SingleAnnotator(
                seed=self.seed,
                base_dir=base_dir,
                annotation_column=self.annotation_key,
                **annotator_config,
            )
            for name, annotator_config in annotators_config.items()
        }

    def _add_missing_primary_keys_(self, df: pd.DataFrame):
        missing_primary_keys = [c for c in self.primary_keys if c not in df.columns]
        if self.is_raise_if_missing_primary_keys:
            if len(missing_primary_keys) > 0:
                raise ValueError(f"Missing primary keys: {missing_primary_keys}")
        else:
            for c in missing_primary_keys:
                df[c] = None

    def _preprocess(self, to_annotate: utils.AnyData) -> pd.DataFrame:
        """Preprocess the examples to annotate. In particular takes care of filtering unnecessary examples."""

        df_to_annotate = utils.convert_to_dataframe(to_annotate)
        self._add_missing_primary_keys_(df_to_annotate)

        for c in self.other_keys_to_keep + [self.annotation_key]:
            if c in df_to_annotate.columns:
                logging.warning(f"{c} column is already in the dataframe. We will overwrite it.")
                df_to_annotate[c] = None

        # remove duplicates because you only need to annotate one of them
        df_to_annotate = df_to_annotate.drop_duplicates(subset=self.primary_keys)

        # set the annotater for each example
        df_to_annotate[self.annotator_column] = df_to_annotate.apply(
            lambda x: utils.random_seeded_choice(
                # we add "annotator" at the beginning to not use the same seed for all tasks
                seed="annotator" + "".join(x[self.random_seed_key]) + str(self.seed),
                choices=list(self.annotators.keys()),
            ),
            axis=1,
        )

        if self.is_avoid_reannotations:
            df_to_annotate = self._apply_cached_annotations(df_to_annotate)

        return df_to_annotate

    def _annotate(self, df_to_annotate: pd.DataFrame, **decoding_kwargs) -> pd.DataFrame:
        """Annotate the examples."""

        df_annotated = df_to_annotate
        for annotator in self.annotators.keys():
            # only annotate examples that have not been annotated yet
            curr_idcs = df_annotated[self.annotator_column] == annotator
            if self.annotation_key in df_annotated.columns:
                curr_idcs &= df_annotated[self.annotation_key].isna()

            logging.info(f"Annotating {curr_idcs.sum()} examples with {annotator}")

            # actual annotation
            curr_annotated = self.annotators[annotator](
                df_annotated.loc[curr_idcs, self.available_fields_to_format],
                **decoding_kwargs,
            )

            df_annotated = self._merge_annotations(df_annotated, curr_annotated)

        return df_annotated

    def _postprocess_and_store_(
        self,
        df_annotated: pd.DataFrame,
        to_annotate: utils.AnyData,
    ) -> list[dict[str, Any]]:
        """Convert the dataframe into a list of dictionaries to be returned, and store current anntations."""

        df_to_annotate = utils.convert_to_dataframe(to_annotate)
        self._add_missing_primary_keys_(df_to_annotate)

        # select available annotations
        if self.is_store_missing_annotations:
            df_annotated[self.annotation_key] = df_annotated[self.annotation_key].fillna(self.TMP_MISSING_ANNOTATION)
        else:
            df_annotated[self.annotation_key] = df_annotated[self.annotation_key].replace(
                self.TMP_MISSING_ANNOTATION, None
            )

        df_annotated = df_annotated[~df_annotated[self.annotation_key].isna()].copy()

        # try converting to int now that no nan. Note this will only do so if possible
        df_annotated[self.annotation_key] = df_annotated[self.annotation_key].astype(self.annotation_type)

        df_annotated = self._filter_annotations_before_storing(df_annotated)
        self._store_annotations_(df_annotated)

        if self.is_store_missing_annotations:
            # put back None
            df_annotated[self.annotation_key] = df_annotated[self.annotation_key].replace(
                self.TMP_MISSING_ANNOTATION, None
            )

        # need to merge with df_to_annotate in case you dropped duplicates
        on = list(self.primary_keys)
        # keeps columns from both df_to_annotate and df_annotated that are useful
        df_annotated = df_annotated[
            self._get_all_keys_to_keep(list(df_to_annotate.columns) + list(df_annotated.columns))
        ]
        df_to_annotate = df_to_annotate[[c for c in df_to_annotate.columns if c not in df_annotated.columns or c in on]]
        # need to remove all other columns before merging if not you will
        df_annotated = df_to_annotate.merge(df_annotated, on=on, how="outer")

        annotated = df_annotated.to_dict(orient="records")

        return annotated

    def _filter_annotations_before_storing(self, df_annotated: pd.DataFrame) -> pd.DataFrame:
        """Filter annotations before storing them."""
        df_annotated = df_annotated[self._get_all_keys_to_keep(df_annotated.columns)]
        return df_annotated

    def _get_all_keys_to_keep(self, current_columns: Sequence) -> list[str]:
        other_keys_to_keep = [c for c in self.other_keys_to_keep if c in current_columns]
        all_keys_to_keep = self.all_keys + [self.annotation_key] + other_keys_to_keep
        return all_keys_to_keep

    def _apply_cached_annotations(self, df_to_annotate: pd.DataFrame) -> pd.DataFrame:
        """annotate examples with cached annotations"""
        df_to_annotate = self._merge_annotations(df_to_annotate, self.df_annotations)
        return df_to_annotate

    def _store_annotations_(self, df_annotated: pd.DataFrame):
        """Store annotation in memory and on disk"""

        if self.df_annotations is None:
            df_annotations = df_annotated
        else:
            df_annotations = pd.concat([self.df_annotations, df_annotated], axis=0, ignore_index=True)

        self.df_annotations = df_annotations.drop_duplicates(subset=self.all_keys, keep="last")

    def _merge_annotations(self, df_to_annotate: pd.DataFrame, df_partially_annotated: pd.DataFrame) -> pd.DataFrame:
        """Merge (partial) annotations with the original df to keep the same order and avoid duplicates annotations."""

        if df_partially_annotated is None or df_partially_annotated.empty:
            return df_to_annotate

        other_keys_to_keep = [c for c in self.other_keys_to_keep if c in df_partially_annotated.columns]

        kwargs = dict(
            on=self.all_keys,
            how="left",
            suffixes=("_old", "_new"),
        )
        try:
            df_to_annotate = df_to_annotate.merge(
                df_partially_annotated[self.all_keys + [self.annotation_key] + other_keys_to_keep],
                **kwargs,
            )
        except ValueError:
            # can have merging issues if columns have different dtypes
            df_partially_annotated = df_partially_annotated.astype({k: str for k in self.all_keys})
            df_to_annotate = df_to_annotate.astype({k: str for k in self.all_keys}).merge(
                df_partially_annotated[self.all_keys + [self.annotation_key] + other_keys_to_keep],
                **kwargs,
            )

        # if columns were in both dataframes, try to merge them
        for c in other_keys_to_keep + [self.annotation_key]:
            if f"{c}_old" in df_to_annotate.columns and f"{c}_new" in df_to_annotate.columns:
                df_to_annotate[c] = df_to_annotate[c + "_old"].fillna(df_to_annotate[c + "_new"])
                df_to_annotate = df_to_annotate.drop(columns=[c + "_old", c + "_new"])

        return df_to_annotate

    #######################


class BaseAnnotatorJSON(BaseAnnotator):
    __doc__ = (
        BaseAnnotator.__doc__.replace(
            "Base class for a pool of annotators.",
            "Base class for a pool of annotators with caching to JSON file.",
        )
        + """
    caching_path : Path, optional
        Path to cache the annotations to. If None, will not save the annotations. If the path already exists it will
        load annotations from there.
    """
    )

    def __init__(self, *args, caching_path: Optional[utils.AnyPath] = "auto", **kwargs):
        super().__init__(*args, **kwargs)
        self.caching_path = self._initialize_cache(caching_path)

    def save(self, path: Optional[utils.AnyPath] = None):
        """Save all annotations to json."""

        path = path or self.caching_path
        if path is not None:
            logging.info(f"Saving all annotations to {path}.")
            # to make sure that we don't overwrite the annotations we load again from file (ideally would use a DB)
            self._refresh_annotations_()
            if not self.is_store_missing_annotations:
                self.df_annotations = self.df_annotations[~self.df_annotations[self.annotation_key].isna()]
            self.df_annotations.to_json(path, orient="records", indent=2)

    def load_(self, path: Optional[utils.AnyPath] = None):
        """Load all the annotations from json."""
        path = path or self.caching_path
        if path is not None:
            path = Path(path)
            if path.exists():
                logging.info(f"Loading all annotations from {path}.")
                self.df_annotations = pd.read_json(path, dtype={k: str for k in self.all_keys})

    def _initialize_cache(self, caching_path):
        if caching_path == "auto":
            if isinstance(self.annotators_config, (str, Path, os.PathLike)):
                stem = Path(self.annotators_config).stem
                caching_path = Path(self.annotators_config).parent / f"annotations_seed{self.seed}_{stem}.json"
                logging.info(f"Saving annotations to `{caching_path}`.")
            else:
                logging.warning("caching_path cannot be 'auto' if annotators_config is not a path. Setting to None.")
                caching_path = None
        elif caching_path is not None:
            logging.warning("Saving_path is given but not 'auto', make sure that it's different for different seeds.")
        self.load_(caching_path)
        return caching_path

    def _store_annotations_(self, df_annotated_to_store: pd.DataFrame):
        super()._store_annotations_(df_annotated_to_store)
        self.save()

    def _refresh_annotations_(self):
        """Refresh the annotations in memory."""
        curr_df_annotations = self.df_annotations.copy()
        self.load_()
        self.df_annotations = pd.concat(
            [self.df_annotations, curr_df_annotations], axis=0, ignore_index=True
        ).drop_duplicates(subset=self.all_keys, keep="last")


class SingleAnnotator:
    """A helper class for a single auto annotator.
    Parameters
    ----------
    prompt_template : str or path
        A prompt template that will be given to `fn_prompter` or path to those prompts. Path is relative to
        `evaluators_configs/`
    fn_completion_parser : callable or str
        Function that maps (parses) the completion to a list of annotations. If a string, it should be a function in
        `completion_parsers.py` to use for parsing the completions into annotations. For each completion, the number of
        annotations (lenght of list) should be equal to the batch_size if not we set all the annotations in that batch
        to NaN.
    completion_parser_kwargs : dict
        Kwargs for fn_completion_parser.
    fn_completions : callable or str
        Function in `decoders.py` to use for decoding the output.
    completions_kwargs : dict
        kwargs for fn_completions. E.g. model_name, max_tokens, temperature, top_p, top_k, stop_seq.
    is_shuffle : bool
        Whether to shuffle the order of the examples before making the prompt. Useful if batch_size > 1.
    seed : int
        Seed for randomization.
    batch_size : int
        Number of examples that will be added in a single prompt.
    base_dir : Path, optional
        Path to the directory containing the annotators configs. I.e. annotators_config will be relative
        to this directory.
    annotation_column : str, optional
        Name of the annotation column in the output dataframe.
    is_store_raw_completions : bool, optional
        Whether to store raw completions at `"raw_completion"` column in the output dataframe.
    processors_to_kwargs : Sequence[dict(str, dict)], optional
        A dictionary of BaseProcessor objects to apply for preprocessing the  dataframe before making the prompts and
        prostprocessing after anntoations. The key should be the names of the BaseProcessor objectsto use in
        `processors.py` the values are the kwargs for the constructor of the Processor. Order matters.
    is_add_default_processors : bool, optional
        Whether to add the default processors to the list of processors.
    """

    def __init__(
        self,
        prompt_template: utils.AnyPath,
        fn_completion_parser: Optional[Union[Callable, str]] = "regex_parser",
        completion_parser_kwargs: Optional[dict[str, Any]] = None,
        fn_completions: Union[Callable, str] = "openai_completions",
        completions_kwargs: Optional[dict[str, Any]] = None,
        is_shuffle: bool = True,
        seed: Optional[int] = 123,
        batch_size: int = 1,
        base_dir: utils.AnyPath = constants.EVALUATORS_CONFIG_DIR,
        annotation_column: str = "annotation",
        is_store_raw_completions: bool = False,
        processors_to_kwargs: Optional[dict[str, dict]] = None,
        is_add_default_processors: bool = True,
    ):
        self.base_dir = Path(base_dir)
        self.prompt_template = self._get_prompt_template(prompt_template)

        if fn_completion_parser is None:
            fn_completion_parser = lambda x: [x]
        elif isinstance(fn_completion_parser, str):
            fn_completion_parser = self._search_fn_completion_parser(fn_completion_parser)
        completion_parser_kwargs = completion_parser_kwargs or {}
        self.fn_completion_parser = partial(fn_completion_parser, **completion_parser_kwargs)

        self.fn_completions = get_fn_completions(fn_completions)
        self.completions_kwargs = completions_kwargs or {}
        self.seed = seed
        self.is_shuffle = is_shuffle
        self.batch_size = batch_size
        self.annotation_column = annotation_column
        self.completion_column = "raw_completion" if is_store_raw_completions else None

        self.is_add_default_processors = is_add_default_processors
        self.processors = []
        processors_to_kwargs = processors_to_kwargs or {}
        if batch_size > 1 and self.is_add_default_processors:
            processors_to_kwargs["PaddingForBatchesProcessor"] = {
                "batch_size": batch_size,
                "padding_example": DUMMY_EXAMPLE,
            }
        for processor, processor_kwargs in processors_to_kwargs.items():
            processor_kwargs["seed"] = self.seed
            Processor = self._search_processor(processor)
            self.processors += [Processor(**processor_kwargs)]

    ### Public methods ###
    def __call__(self, df_to_annotate: pd.DataFrame, **decoding_kwargs) -> pd.DataFrame:
        """Annotates the given examples.
        Parameters
        ----------
        df_to_annotate : pd.DataFrame
            Examples to annotate
        decoding_kwargs :
            Additional arguments to pass to `fn_completions`.
        """
        df_to_annotate = df_to_annotate.copy()  # avoid in place modifications

        if df_to_annotate.empty:
            df_to_annotate[self.annotation_column] = []
            return df_to_annotate

        df_to_annotate = self._preprocess(df_to_annotate)

        # prompts and completions here will not be the same length as the dataframe due to batching
        prompts, df_to_annotate = self._make_prompts(df_to_annotate)

        completions = self.fn_completions(prompts=prompts, **self.completions_kwargs, **decoding_kwargs)

        annotations_to_save, completions_to_save = self._parse_completions(completions=completions["completions"])
        df_to_annotate[self.annotation_column] = annotations_to_save
        if self.completion_column is not None:
            df_to_annotate[self.completion_column] = completions_to_save

        for k, v in completions.items():
            if k != "completions":
                if len(df_to_annotate[self.annotation_column]) == len(v) * self.batch_size:
                    v = [el for el in v for _ in range(self.batch_size)]
                df_to_annotate[k] = v
                if "per_example" in k:
                    df_to_annotate[k] = df_to_annotate[k] / self.batch_size

        df_annotated = self._postprocess(df_to_annotate)

        return df_annotated

    ######################

    ### Private methods ###
    def _search_fn_completion_parser(self, name: str) -> Callable:
        """Search for a completion parser by name."""
        return utils.get_module_attribute(completion_parsers, name)

    def _search_processor(self, name: str) -> Type["processors.BaseProcessor"]:
        """Search for a Processor class by name."""
        return utils.get_module_attribute(processors, name)

    def _get_prompt_template(self, prompt_template: utils.AnyPath):
        return utils.read_or_return(self.base_dir / prompt_template)

    def _make_prompts(
        self, df_to_annotate: pd.DataFrame, prompt_template: Optional[str] = None
    ) -> tuple[list[str], pd.DataFrame]:
        """Make all the prompts for the given examples.
        Parameters
        ----------
        df_to_annotate : pd.DataFrame
            Examples to annotate
        prompt_template : str
            Template to use for the prompt. If None, use the one from the constructor.
        Returns
        -------
        prompts : list[str]
            Formatted prompts for the given examples.
        df_to_annotate : pd.DataFrame
            Examples to annotate in the same order as prompts.
        """
        if prompt_template is None:
            prompt_template = self.prompt_template
        return utils.make_prompts(df=df_to_annotate, template=prompt_template, batch_size=self.batch_size)

    def _preprocess(self, df_to_annotate: pd.DataFrame) -> pd.DataFrame:
        """Preprocess the examples before annotating. In particular, takes care of all the randomization."""

        for processor in self.processors:
            df_to_annotate = processor.preprocess(df_to_annotate)

        if self.is_shuffle:
            df_to_annotate = df_to_annotate.sample(frac=1, random_state=self.seed)

        return df_to_annotate

    def _parse_completions(self, completions: list[str]) -> tuple[list[Any], list[Any]]:
        """Converts the completions into annotations."""
        all_annotations = []
        all_completions = []
        for completion in completions:
            try:
                batch_annotations = self.fn_completion_parser(completion)
                batch_annotations = list(batch_annotations)

                if len(batch_annotations) != self.batch_size:
                    logging.warning(
                        f"Found {len(batch_annotations)} annotations in:'''\n{completion}\n''' but expected"
                        f" {self.batch_size}. We are setting all annotations to None."
                    )
                    batch_annotations = [None] * self.batch_size

            except Exception as e:
                logging.exception(f"Error while parsing completion: '''\n{completion}\n'''")
                batch_annotations = [None] * self.batch_size

            all_annotations += batch_annotations
            all_completions += [completion] * self.batch_size
        return all_annotations, all_completions

    def _postprocess(self, df_annotated: pd.DataFrame) -> pd.DataFrame:
        """Postprocess the annotated examples."""

        arr_is_na = df_annotated[self.annotation_column].isna()
        if arr_is_na.any():
            logging.warning(
                f"{arr_is_na.sum().item()} samples had no auto annotation. We are filtering them for now. "
                f"If you are using chain of thought it might be that max_tokens limit is too low. "
            )
            df_annotated = df_annotated[~arr_is_na]

        for processor in self.processors[::-1]:  # postprocess in reverted order => no interactions between processors
            df_annotated = processor.postprocess(df_annotated)

        return df_annotated

    #######################