# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

# Assume saving data from https://github.com/tonyzhaozh/few-shot-learning/tree/main/data/slot-movies in the current folder

import pandas as pd
import json
import pickle
import numpy as np
import os


# Imported from https://github.com/tonyzhaozh/few-shot-learning/blob/main/data_utils.py#L122-#L173
def generate_dataset(field_name, data_path="."):

    all_fields = ["Actor", "Award", "Character_Name", "Director", "Genre", "Opinion", "Origin", "Plot", "Quote", "Relationship", "Soundtrack", "Year"]
    assert field_name in all_fields
    all_fields.remove(field_name)
    filter_tags = [f"B-{field}" for field in all_fields] + [f"I-{field}" for field in all_fields] + ["O"]
    target_tags = [f"B-{field_name}", f"I-{field_name}"]

    with open(f'{data_path}/slot-movies/train', 'r') as f:
        lines = f.readlines()
        lines = [line.replace(' <=> <NULL>','').strip() for line in lines]
    train_answers = []
    train_sentences = []
    for line in lines:
        answer = ''
        untagged_line = ''
        for word in line.split(' '):
            contains_target = [tag in word for tag in target_tags]
            if np.any(contains_target):
                for tag in target_tags:
                    word = word.replace(':' + tag, '')
                answer += word + ' '
            for tag in filter_tags:
                word = word.replace(':' + tag, '')
            untagged_line += word + ' '

        if answer != '':
            train_answers.append(answer.strip())
            train_sentences.append(untagged_line.strip())


    with open(f'{data_path}/slot-movies/test', 'r') as f:
        lines = f.readlines()
        lines = [line.replace(' <=> <NULL>','').strip() for line in lines]
    test_answers = []
    test_sentences = []
    for line in lines:
        answer = ''
        untagged_line = ''
        for word in line.split(' '):
            contains_target = [tag in word for tag in target_tags]
            if np.any(contains_target):
                for tag in target_tags:
                    word = word.replace(':' + tag, '')
                answer += word + ' '
            for tag in filter_tags:
                word = word.replace(':' + tag, '')
            untagged_line += word + ' '

        if answer != '':
            test_answers.append(answer.strip())
            test_sentences.append(untagged_line.strip())
    if not os.path.isdir((f"{data_path}/movie/{field_name}")):
        os.makedirs(f"{data_path}/movie/{field_name}", exist_ok=True)
    train_data = {}
    train_data['content'] = train_sentences
    train_data['label'] = train_answers
    df = pd.DataFrame(train_data)
    df.to_csv(f"{data_path}/movie/{field_name}/train.csv")

    test_data = {}
    test_data['content'] = test_sentences
    test_data['label'] = test_answers
    df = pd.DataFrame(test_data)
    df.to_csv(f"{data_path}/movie/{field_name}/test.csv")

for field_name in ['Director', 'Genre']:
    # by default save to ./movie/field_name, this is consistent with src/run_exp_movie.py#L362 
    generate_dataset(field_name, data_path='./')
data/process_movie.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import asyncio
import contextlib
from collections.abc import AsyncGenerator, Callable, Iterable
from typing import Any, TypeVar

T = TypeVar("T")
U = TypeVar("U")


class FinishedMarker:
    pass


async def bucket(
    num_buckets: int,
    it: AsyncGenerator[T, None],
    process: Callable[[T], Iterable[tuple[int, U]]],
) -> list[AsyncGenerator[U, None]]:
    """
    Wraps `it` and distributes its items into `num_buckets` buckets based on the
    result of `process`.

    Similar to https://more-itertools.readthedocs.io/en/stable/api.html#more_itertools.bucket
    but with the following differences:
    - The buckets are integers in [0, num_buckets).
    - The `process` callable returns both the bucket and a transformed version
      of the element from `it`.
    """
    queues: list[asyncio.Queue[U | FinishedMarker]] = [
        asyncio.Queue() for _ in range(num_buckets)
    ]
    # Keep track of how many consumers are still open, so if they're all closed,
    # `fill_queues` can finish and close `it` early.
    num_open_consumers = num_buckets

    async def fill_queues() -> None:
        try:
            async with contextlib.aclosing(it):
                async for item in it:
                    if num_open_consumers == 0:
                        break
                    for bucket_index, processed_item in process(item):
                        queues[bucket_index].put_nowait(processed_item)
        finally:
            for queue in queues:
                queue.put_nowait(FinishedMarker())

    fill_queues_task = asyncio.create_task(fill_queues())

    async def gen(queue: asyncio.Queue[U | FinishedMarker]) -> AsyncGenerator[U, None]:
        nonlocal num_open_consumers
        try:
            # Below, we will run the generator for one iteration for returning it,
            # to ensure that it starts executing.  We would like this so that when
            # the generator is closed, it runs the finally block below to decrement
            # `num_open_consumers`. If the generator is closed before it has run any
            # iterations, then the finalizer below will not run, and
            # `num_open_consumers` will not be decremented.
            # Therefore, we yield nothing once here.
            yield  # type: ignore
            while True:
                next_item = await queue.get()
                if isinstance(next_item, FinishedMarker):
                    break
                yield next_item

            # Check that there were no exceptions in the fill_queues_task
            # TODO(richard): If there was an exception, the printed traceback isn't very useful; improve that.
            await fill_queues_task
        finally:
            num_open_consumers -= 1

    # Run each generator for one iteration.
    result = [gen(queue) for queue in queues]
    for result_item in result:
        await anext(result_item)
    return result


async def consume(it: AsyncGenerator[Any, None]) -> None:
    """Fully executes the async generator `it` while discarding all output."""
    async with contextlib.aclosing(it):
        async for _ in it:
            pass

lmapi/src/lmapi/async_tools/asyncitertools.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import asyncio
import collections
import dataclasses
import functools
import time
from collections.abc import Awaitable, Callable
from contextlib import AbstractAsyncContextManager, asynccontextmanager
from dataclasses import dataclass
from typing import Any, TypeVar, cast

_MICROS_RATIO = 1_000_000
_NANOS_RATIO = 1_000_000_000


@dataclass
class TokenBucket:
    """Used to prevent too much resource consumption during any time period.

    These tokens are not words (like in NLP) and the bucket is not a map.
    Instead the tokens are like coins and the bucket is like a physical bucket.

    See https://en.wikipedia.org/wiki/Token_bucket#Algorithm for more details.
    However, unlike typical implementations of token buckets, this one allows
    the refill rate to change at any time.
    """

    # Maximum number of tokens that can be in the bucket, multiplied by 1_000_000.
    capacity_micros: int
    # Amount of tokens added per second, multipled by 1_000_000.
    refill_rate_micros: int

    # The current amount of tokens in the bucket.
    level_micros: int = 0

    # Time when `level` was last updated, in nanoseconds.
    last_updated_ns: int = dataclasses.field(default_factory=time.monotonic_ns)
    # One entry per call to `claim` that's waiting for the bucket to refill.
    _waiting: collections.deque[tuple[asyncio.Event, int]] = dataclasses.field(
        default_factory=collections.deque
    )
    # Set by the next `claim` in line so that it can be notified if the refill rate changes.
    _next_token_refill: asyncio.Event | None = None

    @staticmethod
    def create(capacity: float, refill_rate: float, initial_level: float = 0):
        assert capacity >= 0
        assert refill_rate >= 0
        assert initial_level >= 0

        return TokenBucket(
            int(capacity * _MICROS_RATIO),
            int(refill_rate * _MICROS_RATIO),
            int(initial_level * _MICROS_RATIO),
        )

    async def claim(self, tokens: float) -> None:
        tokens_micros = int(tokens * _MICROS_RATIO)
        assert tokens_micros <= self.capacity_micros
        del tokens

        event: asyncio.Event | None = None
        try:
            # Check if we have sufficient tokens already
            self._update_level()
            # async with self._lock:
            if self.level_micros >= tokens_micros:
                self.level_micros -= tokens_micros
                return

            event = asyncio.Event()
            self._waiting.append((event, tokens_micros))
            while True:
                # Check whether we're first in line
                first_event, _ = self._waiting[0]
                if event is first_event:
                    # We might already have enough tokens. If so, don't wait for more tokens to fill.
                    if self.level_micros >= tokens_micros:
                        break
                    try:
                        # Wait for the bucket to fill up.
                        self._next_token_refill = event

                        # Compute number of microseconds to wait until bucket will be full,
                        # if the rate doesn't change in the meantime.
                        # Round up to the next microsecond.
                        timeout_micros, remainder = divmod(
                            (tokens_micros - self.level_micros) * _MICROS_RATIO,
                            self.refill_rate_micros,
                        )
                        if remainder:
                            timeout_micros += 1

                        # Sleep for the computed number of microseconds.
                        # We assume that 1) the floating point value for timeout
                        # has sufficient precision to represent the desired
                        # number of microseconds, and 2) wait_for's timer has
                        # sufficient resolution.
                        # If the rate changes before the timeout, the event will fire.
                        await asyncio.wait_for(
                            event.wait(), timeout=timeout_micros / _MICROS_RATIO
                        )

                        # If we reach here, before we refilled all the tokens,
                        # the refill rate changed. Recompute how long we should
                        # wait until the tokens are refilled, and try again.
                        # No matter how the rate changed, we shouldn't have
                        # enough tokens at this point.
                        event.clear()
                        self._next_token_refill = None
                        continue
                    except asyncio.TimeoutError:
                        # We should have collected enough tokens by now to take our turn
                        self._update_level()
                        break
                else:
                    # Wait until we get to the start of the line
                    await event.wait()
                    event.clear()

            # async with self._lock:
            assert self.level_micros >= tokens_micros, (
                self.level_micros,
                tokens_micros,
            )
            self.level_micros -= tokens_micros

            self._waiting.popleft()
            if self._waiting:
                # Tell the next in line that they're first now.
                event, _ = self._waiting[0]
                event.set()
            return

        except asyncio.CancelledError:
            pass

    def add_to_rate(
        self, delta: float, min_rate: float | None = None, max_rate: float | None = None
    ) -> None:
        """Add to the refill rate.

        Supply a negative number to subtract from the rate. However, the
        negative number cannot cause the rate to become negative."""
        new_rate = self.refill_rate_micros + round(delta * _MICROS_RATIO)
        assert new_rate >= 0, "Cannot have a negative refill rate"
        if min_rate is not None:
            new_rate = max(new_rate, round(min_rate * _MICROS_RATIO))
        if max_rate is not None:
            new_rate = min(new_rate, round(max_rate * _MICROS_RATIO))
        self.refill_rate_micros = new_rate
        self.reset_rate(new_rate)

    def multiply_rate(
        self, ratio: float, min_rate: float | None = None, max_rate: float | None = None
    ) -> None:
        """Apply a multiplicative factor to the refill rate."""

        assert ratio >= 0, "Cannot have a negative ratio"
        new_rate = round(self.refill_rate_micros * ratio)
        if min_rate is not None:
            new_rate = max(new_rate, round(min_rate * _MICROS_RATIO))
        if max_rate is not None:
            new_rate = min(new_rate, round(max_rate * _MICROS_RATIO))
        self.reset_rate(new_rate)

    def reset_rate(self, new_rate: float) -> None:
        """Reset the refill rate."""
        # Refill tokens with the previous rate, before updating the rate.
        self._update_level()

        self.refill_rate_micros = round(new_rate * _MICROS_RATIO)
        # Tell the task waiting for tokens to refill to wake up so that we can
        # recompute how much longer it should sleep, with the current rate.
        if self._next_token_refill:
            self._next_token_refill.set()

    def _update_level(self) -> None:
        now = time.monotonic_ns()
        elapsed_nanos = now - self.last_updated_ns
        self.level_micros = min(
            self.capacity_micros,
            self.level_micros
            + (elapsed_nanos * self.refill_rate_micros) // _NANOS_RATIO,
        )
        self.last_updated_ns = now


class RateLimitExceededError(Exception):
    pass


AsyncCallableT = TypeVar("AsyncCallableT", bound=Callable[..., Awaitable[Any]])
AsyncContextManagerProducerT = TypeVar(
    "AsyncContextManagerProducerT",
    bound=Callable[..., AbstractAsyncContextManager[Any]],
)


@dataclass
class AdaptiveLimiter:
    """Adaptively rate-limit the wrapped function.

    When a call to the function is successful, we increase the rate additively;
    if it complains that the rate was exceeded, we decrease the rate multiplicatively.
    """

    initial_qps: float = 10
    max_qps: float = 500
    min_qps: float = 1
    bucket: TokenBucket = dataclasses.field(init=False)

    def __post_init__(self):
        self.bucket = TokenBucket.create(self.max_qps, self.initial_qps)

    def wrap_async_callable(self, func: AsyncCallableT) -> AsyncCallableT:
        @functools.wraps(func)
        async def wrapped(*args, **kwargs):
            while True:
                # Wait our turn
                await self.bucket.claim(1)
                # Try calling the function
                try:
                    result = await func(*args, **kwargs)
                    self.bucket.add_to_rate(1, max_rate=self.max_qps)
                    return result
                except RateLimitExceededError:
                    self.bucket.multiply_rate(0.9, min_rate=self.min_qps)

        return cast(AsyncCallableT, wrapped)

    def wrap_async_context_manager_producer(
        self, func: AsyncContextManagerProducerT
    ) -> AsyncContextManagerProducerT:
        @functools.wraps(func)
        @asynccontextmanager
        async def wrapped(*args, **kwargs):
            while True:
                # Wait our turn
                await self.bucket.claim(1)
                # Try calling the function
                try:
                    async with func(*args, **kwargs) as cm_result:
                        yield cm_result
                    self.bucket.add_to_rate(1, max_rate=self.max_qps)
                    break
                except RateLimitExceededError:
                    self.bucket.multiply_rate(0.9, min_rate=self.min_qps)

        return cast(AsyncContextManagerProducerT, wrapped)

lmapi/src/lmapi/async_tools/limits.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from collections.abc import AsyncIterable, AsyncIterator
from dataclasses import dataclass


@dataclass
class ServerSentEvent:
    event: str | None = None
    data: str | None = None
    id: str | None = None
    retry: int | None = None


async def parse_event_stream(
    lines: AsyncIterable[str],
) -> AsyncIterator[ServerSentEvent]:
    """Implements part of https://html.spec.whatwg.org/multipage/server-sent-events.html#event-stream-interpretation
    This function doesn't implement the logic specified there for dispatching the event; it only performs parsing.
    """
    next_event = ServerSentEvent()

    async for line in lines:
        line = line.rstrip("\n")
        if len(line) == 0:
            yield next_event
            next_event = ServerSentEvent()
            continue

        if line[0] == ":":
            continue

        field, _, value = line.partition(":")
        if value[0] == " ":
            value = value[1:]
        if field == "event":
            next_event.event = value
        elif field == "data":
            if next_event.data is None:
                next_event.data = value + "\n"
            else:
                next_event.data += value + "\n"
        elif field == "id":
            next_event.id = value
        elif field == "retry":
            try:
                next_event.retry = int(value)
            except ValueError:
                pass

    # No need to handle `next_event` now, as the specification states:
    #   Once the end of the file is reached, any pending data must be discarded.
    #   (If the file ends in the middle of an event, before the final empty line,
    #   the incomplete event is not dispatched.)

lmapi/src/lmapi/async_tools/server_sent_events.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import dataclasses
from dataclasses import InitVar, dataclass
from typing import Protocol


class AuthorizationProvider(Protocol):
    def headers(self) -> dict[str, str]:
        """Returns the headers to be used for authorization."""
        ...


@dataclass
class OpenAiApiKey(AuthorizationProvider):
    key: InitVar[str]
    _headers: dict[str, str] = dataclasses.field(init=False)

    def __post_init__(self, key: str) -> None:
        self._headers = {"Authorization": f"Bearer {key}"}

    def headers(self) -> dict[str, str]:
        return self._headers


@dataclass
class AoaiApiKey(AuthorizationProvider):
    key: InitVar[str]
    _headers: dict[str, str] = dataclasses.field(init=False)

    def __post_init__(self, key: str) -> None:
        self._headers = {"api-key": key}

    def headers(self) -> dict[str, str]:
        return self._headers

lmapi/src/lmapi/auth.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from abc import abstractmethod
from collections.abc import AsyncGenerator, Mapping, Sequence
from dataclasses import dataclass
from typing import Protocol, TypeAlias

import tiktoken


class TokenWithLogprob(Protocol):
    @property
    def text(self) -> str:
        """The textual representation of the token.

        For tokens that don't constitute valid UTF-8, this will look like "bytes:\\x??\\x??".
        """
        ...

    @property
    def bytes(self) -> bytes:
        "The bytes that make up the token."
        ...

    @property
    def logprob(self) -> float:
        ...

    @property
    def token_id(self) -> int:
        ...


@dataclass
class SampledToken:
    token: TokenWithLogprob
    top_choices: tuple[TokenWithLogprob, ...]


# Types:
# - rules for how we can interrupt the completion: stop, etc.
# - whether we can get logprobs

Completion: TypeAlias = Sequence[SampledToken]


@dataclass
class CompletionsSettings:
    """The meanings are the same as in https://learn.microsoft.com/en-us/azure/cognitive-services/openai/reference#completions

    If any field is set to None, the defaults from that documentation page apply."""

    temperature: float | None = None
    max_tokens: int | None = None
    n: int | None = None
    stop: str | Sequence[str] | None = None
    logprobs: int | None = None
    echo: bool | None = None
    logit_bias: Mapping[int, float] | None = None
    top_p: float | None=None


class LM(Protocol):
    @abstractmethod
    async def completions(
        self,
        prompt: str | Sequence[int] | Sequence[str] | Sequence[Sequence[int]],
        settings: CompletionsSettings | None = None,
    ) -> Sequence[Completion]:
        """Returns completions given the prompt."""
        ...

    @abstractmethod
    async def streaming_completions(
        self, prompt: str | Sequence[int], settings: CompletionsSettings | None = None
    ) -> Sequence[AsyncGenerator[SampledToken, None]]:
        """Returns completions in a streaming fashion as async generators of SampledText.

        To ensure that the HTTP connection to GPT is closed quickly, use `contextlib.aclosing`:
        ```
        # streams is a sequence with length equal to `n` in the settings,
        # the number of completions to generate given the prompt.
        # In this example, we assume n = 1.
        streams = gpt.streaming_completions(prompt)
        [stream] = streams
        async with contextlib.aclosing(stream):
            async for sampled_text in stream:
                ...
                if some_condition:
                    break
        ```

        This pattern ensures that when the loop exits with the `break`, `it.aclose()` is called immediately,
        rather than only when `it` is garbage collected.
        """
        ...

    @property
    def encoding(self) -> tiktoken.Encoding:
        """Returns the encoding scheme the model uses."""
        ...

lmapi/src/lmapi/lm.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import ast
import dataclasses
import json
from collections.abc import (
    AsyncGenerator,
    AsyncIterable,
    AsyncIterator,
    Callable,
    Iterator,
    Sequence,
)
from contextlib import AbstractAsyncContextManager, asynccontextmanager
from contextvars import ContextVar
from dataclasses import dataclass
from functools import cached_property
from typing import TYPE_CHECKING, Any, Protocol, TypeVar

import aiohttp
import tiktoken

from lmapi.async_tools import asyncitertools, limits
from lmapi.async_tools.server_sent_events import ServerSentEvent, parse_event_stream
from lmapi.auth import AuthorizationProvider
from lmapi.lm import LM, Completion, CompletionsSettings, SampledToken, TokenWithLogprob

client_session: ContextVar[aiohttp.ClientSession] = ContextVar("client_session")


@dataclass(frozen=True)
class OpenAIAPIError(Exception):
    """Indicates a general OpenAI call error."""

    status_code: int
    text: str

    @property
    def user_message(self) -> str:
        return "Model communication failure"

    @cached_property
    def debug_message(self) -> str:
        return f"Unexpected status code: {self.status_code}. {self.text}"


@dataclass(slots=True)
class OpenAITokenWithLogprob:
    if TYPE_CHECKING:

        def _check_protocol(self) -> TokenWithLogprob:
            return self

    text: str
    logprob: float
    _encoding: tiktoken.Encoding

    _bytes: bytes | None = dataclasses.field(init=False, default=None)

    @property
    def bytes(self) -> bytes:
        if self._bytes is None:
            self._bytes = openai_token_to_bytes(self.text)
        return self._bytes

    _token_id: int | None = dataclasses.field(init=False, default=None)

    @property
    def token_id(self) -> int:
        if self._token_id is None:
            self._token_id = self._encoding.encode_single_token(self.bytes)
        return self._token_id


class _NextLogprobsFunction(Protocol):
    async def __call__(self, prompt: str | Sequence[int]) -> dict[int, float]:
        ...


class _CompletionsFunction(Protocol):
    async def __call__(
        self,
        prompt: str | Sequence[int] | Sequence[str] | Sequence[Sequence[int]],
        settings: CompletionsSettings | None = None,
    ) -> Sequence[Completion]:
        ...


@dataclass(frozen=True)
class OpenAI:
    """Implementation of the LM protocol for OpenAI models.

    Arguments:
        url: The URL of the model.
        auth_provider: The authorization provider to use.
        encoding: The encoding of the model.
        default_completion_settings: The default completion settings to use.
    """

    if TYPE_CHECKING:

        def _check_protocol(self) -> LM:
            return self

    url: str
    auth_provider: AuthorizationProvider
    encoding: tiktoken.Encoding
    default_completion_settings: dict[str, Any]
    additional_headers: dict[str, str]
    request_limiter: limits.AdaptiveLimiter | None = None

    @staticmethod
    def create(
        url: str,
        encoding_or_name: str | tiktoken.Encoding,
        auth_provider: AuthorizationProvider,
        request_limiter: limits.AdaptiveLimiter | None = None,
        default_completion_settings: dict[str, Any] | None = None,
        additional_headers: dict[str, Any] | None = None,
    ) -> "OpenAI":
        encoding = (
            tiktoken.get_encoding(encoding_or_name)
            if isinstance(encoding_or_name, str)
            else encoding_or_name
        )

        return OpenAI(
            url,
            auth_provider,
            encoding,
            default_completion_settings or {},
            additional_headers or {},
            request_limiter,
        )

    async def completions(
        self,
        prompt: str | Sequence[int] | Sequence[str] | Sequence[Sequence[int]],
        settings: CompletionsSettings | None = None,
    ) -> Sequence[Completion]:
        return await self._completions_maybe_limited(prompt, settings)

    @cached_property
    def _completions_maybe_limited(self) -> _CompletionsFunction:
        if self.request_limiter is None:
            return self._completions
        return self.request_limiter.wrap_async_callable(self._completions)

    async def _completions(
        self,
        prompt: str | Sequence[int] | Sequence[str] | Sequence[Sequence[int]],
        settings: CompletionsSettings | None = None,
    ) -> Sequence[Completion]:
        """
        Inner implementation of `completions` before `self.request_limiter` is applied.

        Must be called in an `async with client_session:` block where
        `client_session` is the same one used to construct this object.
        """
        params = self._make_params(prompt, settings)
        async with client_session.get().post(
            self.url,
            headers={**self.auth_provider.headers(), **self.additional_headers},
            json=params,
        ) as response:
            if response.status != 200:
                if response.status in (408, 429, 500):
                    raise limits.RateLimitExceededError()
                else:
                    raise OpenAIAPIError(response.status, await response.text())
            resp = await response.json(content_type=None)
            if resp is None:
                raise limits.RateLimitExceededError()
        result = [
            extract_sampled_tokens(choice["logprobs"], self.encoding)
            for choice in resp["choices"]
        ]
        return result

    async def streaming_completions(
        self, prompt: str | Sequence[int], settings: CompletionsSettings | None = None
    ) -> Sequence[AsyncGenerator[SampledToken, None]]:
        """Please see docstring for Gpt.streaming_completions."""

        params = self._make_params(prompt, settings)
        n = params.get("n", 1)

        async def drop_unneeded(
            events: AsyncIterable[ServerSentEvent],
        ) -> AsyncIterator[dict[str, Any]]:
            """Process the stream of ServerSentEvents to drop the ones that we don't need.

            We drop:
            - The [DONE] event, which normally occurs one at the end
            - Unexpected events where there are no tokens sampled
            """
            async for event in events:
                assert event.data is not None
                if event.data == "[DONE]\n":
                    continue

                data = json.loads(event.data)
                assert len(data["choices"]) == 1
                choice = data["choices"][0]

                if len(choice["logprobs"].get("tokens", [])) == 0:
                    # Sometimes the API sends us an event even though no tokens were sampled.
                    # Ignore those cases.
                    continue

                yield choice

        def process_choice(
            choice: dict[str, Any]
        ) -> Iterator[tuple[int, SampledToken]]:
            """Extracts the data returned by OpenAI's API into the SampledText object."""
            # choice["text"] is always a valid Unicode string, and may correspond to multiple tokens.
            # Sometimes, choice["text"] is empty while choice["logprobs"]["tokens"] is not.
            # This seems to happen when choice["logprobs"]["tokens"] doesn't concatenate into valid UTF-8.
            # sampled_text =  SampledText(choice["text"], choice["finish_reason"])
            for sampled_tokens in extract_sampled_tokens(
                choice["logprobs"], self.encoding
            ):
                yield choice["index"], sampled_tokens

        if n == 1:

            async def gen_1() -> AsyncGenerator[SampledToken, None]:
                async with self.streaming_completions_client_response(
                    params
                ) as response:
                    # We use "utf-8" because
                    # https://html.spec.whatwg.org/multipage/server-sent-events.html#event-stream-interpretation says:
                    # > Streams must be decoded using the UTF-8 decode algorithm.
                    events = parse_event_stream(
                        _bytes_to_str(response.content, "utf-8")
                    )
                    # `response.content` is an async iterable which returns sequences of
                    # bytes ending in b'\n'. That is slightly inappropriate because the SSE
                    # specification allows lines to end in b'\r\n' or b'\r' as well.
                    # But in practice, the GPT server never seems to use \r\n or \r.
                    async for c in drop_unneeded(events):
                        for _, item in process_choice(c):
                            yield item

            return [gen_1()]
        else:

            async def gen_n() -> AsyncGenerator[dict[str, Any], None]:
                async with self.streaming_completions_client_response(
                    params
                ) as response:
                    events = parse_event_stream(
                        _bytes_to_str(response.content, "utf-8")
                    )
                    async for event in drop_unneeded(events):
                        yield event

            return await asyncitertools.bucket(n, gen_n(), process_choice)

    @cached_property
    def streaming_completions_client_response(
        self,
    ) -> Callable[
        [dict[str, Any]], AbstractAsyncContextManager[aiohttp.ClientResponse]
    ]:
        """Helper function for `streaming_completions`.

        This method returns `aiohttp.ClientResponse`; the data retrieved from it is parsed by `streaming_completions`.
        """
        if self.request_limiter is None:
            return self._streaming_completions_client_response
        return self.request_limiter.wrap_async_context_manager_producer(
            self._streaming_completions_client_response
        )

    @asynccontextmanager
    async def _streaming_completions_client_response(
        self, params: dict[str, Any]
    ) -> AsyncIterator[aiohttp.ClientResponse]:
        """Call as: `with gpt_impl._streaming_completions_client_response(params) as response: ..."""

        async with client_session.get().post(
            self.url,
            headers={**self.auth_provider.headers(), **self.additional_headers},
            json={**params, "stream": True},
        ) as response:
            if response.status != 200:
                if response.status in (408, 429, 500):
                    raise limits.RateLimitExceededError(
                        response.status, await response.text()
                    )
                else:
                    raise OpenAIAPIError(response.status, await response.text())
            yield response

    def _make_params(
        self,
        prompt: str | Sequence[int] | Sequence[str] | Sequence[Sequence[int]],
        settings: CompletionsSettings | None,
    ) -> dict[str, Any]:
        params = {"prompt": prompt, "logprobs": 0, **self.default_completion_settings}
        if settings is not None:
            params.update(_filter_none_values(dataclasses.asdict(settings)))
        return params


K = TypeVar("K")
V = TypeVar("V")


def _filter_none_values(d: dict[K, V | None]) -> dict[K, V]:
    return {k: v for k, v in d.items() if v is not None}


async def _bytes_to_str(
    bs: AsyncIterable[bytes], encoding: str
) -> AsyncGenerator[str, None]:
    async for b in bs:
        yield b.decode(encoding)


def openai_token_to_bytes(token: str) -> bytes:
    if token.startswith("bytes:"):
        return ast.literal_eval(f"b'{token[6:]}'")
    else:
        return token.encode("utf-8")


def extract_sampled_tokens(
    logprobs_info: dict, encoding: tiktoken.Encoding
) -> list[SampledToken]:
    tokens = logprobs_info["tokens"]
    token_logprobs = logprobs_info["token_logprobs"]
    top_logprobs = logprobs_info.get("top_logprobs")
    if top_logprobs is None:
        top_logprobs = [{}] * len(tokens)

    sampled_tokens: list[SampledToken] = []
    for token, token_logprob, top_logprobs_for_token in zip(
        tokens, token_logprobs, top_logprobs, strict=True
    ):
        if top_logprobs_for_token is None:
            top_logprobs_for_token = {}
        sampled_tokens.append(
            SampledToken(
                OpenAITokenWithLogprob(token, token_logprob, encoding),
                tuple(
                    OpenAITokenWithLogprob(t, lp, encoding)
                    for t, lp in top_logprobs_for_token.items()
                ),
            )
        )

    return sampled_tokens

lmapi/src/lmapi/openai.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import asyncio
import itertools
from collections.abc import AsyncGenerator

import pytest

from lmapi.async_tools.asyncitertools import bucket


@pytest.mark.asyncio
async def test_bucket_normal() -> None:
    async def gen() -> AsyncGenerator[int, None]:
        for i in range(4):
            if i % 2 == 0:
                await asyncio.sleep(0)
            yield i

    it_indices_permutations = set(itertools.permutations(i % 3 for i in range(4)))
    for it_indices in it_indices_permutations:
        bucketed_its = await bucket(3, gen(), lambda x: [(x % 3, x)])
        regular_its = [iter([0, 3]), iter([1]), iter([2])]
        for it_index in it_indices:
            assert await anext(bucketed_its[it_index]) == next(regular_its[it_index])

        for it in bucketed_its:
            with pytest.raises(StopAsyncIteration):
                await anext(it)

        for it in regular_its:
            with pytest.raises(StopIteration):
                next(it)


class ExceptionForTest(Exception):
    pass


@pytest.mark.asyncio
async def test_bucket_exception() -> None:
    async def gen() -> AsyncGenerator[int, None]:
        for i in range(6):
            if i >= 5:
                raise ExceptionForTest

            if i % 2 == 0:
                await asyncio.sleep(0)
            yield i

    it0, it1, it2 = await bucket(3, gen(), lambda x: [(x % 3, x)])

    assert await anext(it0) == 0
    assert await anext(it0) == 3
    with pytest.raises(ExceptionForTest):
        await anext(it0)

    assert await anext(it1) == 1
    assert await anext(it1) == 4
    with pytest.raises(ExceptionForTest):
        await anext(it1)

    assert await anext(it2) == 2
    with pytest.raises(ExceptionForTest):
        await anext(it2)


@pytest.mark.asyncio
async def test_bucket_closing() -> None:
    """Tests that `bucket` closes the generator when all consumers are closed."""

    lock = asyncio.Lock()
    max_i = None
    gen_was_closed = False

    async def gen() -> AsyncGenerator[int, None]:
        nonlocal max_i, gen_was_closed
        try:
            for i in range(9):
                async with lock:
                    max_i = i
                    yield i
                    await asyncio.sleep(0)
        except GeneratorExit:
            gen_was_closed = True

    async with lock:
        it0, it1, it2 = await bucket(3, gen(), lambda x: [(x % 3, x)])
        await it0.aclose()
        await it2.aclose()

    # Give up control so that `gen` executes one iteration, yielding 0
    await asyncio.sleep(0)
    # Give up control so that `gen` executes one iteration, yielding 1
    await asyncio.sleep(0)

    async with lock:
        assert await anext(it1) == 1
        await it1.aclose()

    # `fill_queues` will run more iteration of `gen` and then realize that
    # `num_open_consumers == 0`, thus closing `gen` as well
    await asyncio.sleep(0)
    assert max_i == 2
    assert gen_was_closed

    # There should be no changes due to this additional sleep
    await asyncio.sleep(0)
    assert max_i == 2
    assert gen_was_closed

lmapi/tests/test_lmapi/async_tools/test_asyncitertools.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

# ruff: noqa: F401
def test_package_import():
    "Tests that lmapi is importable."
    import lmapi

lmapi/tests/test_lmapi/test_import.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from prv_accountant.dpsgd import DPSGDAccountant
import numpy as np
from prv_accountant.privacy_random_variables import PureDPMechanism
from prv_accountant import PRVAccountant
import math

###assumming sensitivity as 1.
dataset = "DBPedia"
assert dataset in ["DBPedia"]
sigma_list = []

if dataset == "DBPedia":
    full_train_num = 40000
    n = 80
    max_len_token = 100
    sigma_list = [2.73, 3.34, 3.95, 4.57]


sample_rate = n / full_train_num
print(n, full_train_num)
print("sample rate", sample_rate)
print("steps", max_len_token)

for sigma in sigma_list:
    sigma_bar = math.log(1+sample_rate*(math.exp(sigma)-1))
    print(sigma_bar)
    prv_0 = PureDPMechanism(sigma_bar)

    accountant = PRVAccountant(
        prvs=[prv_0, ],
        max_self_compositions=[max_len_token+1],
        eps_error=0.01,
        delta_error=1e-10
    )
    eps_low, eps_est, eps_up = accountant.compute_epsilon(delta=1/full_train_num, num_self_compositions=[max_len_token])

    print(sigma, eps_low, eps_est, eps_up)

privacy_analysis/exponential.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from prv_accountant.dpsgd import DPSGDAccountant
import numpy as np
from prv_accountant.privacy_random_variables import PoissonSubsampledGaussianMechanism, GaussianMechanism, LaplaceMechanism
from prv_accountant import PRVAccountant


dataset = "MIT-G"
assert dataset in ["AGNEWS", "DBPedia", "MIT-D", "MIT-G"]
sigma_list = []

if dataset == "AGNEWS":
    full_train_num = 30000
    n = 20
    max_token_cnt = 100
    sigma_list = [0.51, 0.46, 0.39, 0.31]
elif dataset == "DBPedia":
    full_train_num = 40000
    n = 80
    max_token_cnt = 100
    sigma_list = [0.63, 0.54, 0.45, 0.36]
elif dataset == "MIT-G":
    full_train_num = 2953
    n = 80
    max_token_cnt = 80
    sigma_list = [1.08, 0.81, 0.64, 0.5]
elif dataset == "MIT-D":
    full_train_num = 1561
    n = 80
    max_token_cnt = 80
    sigma_list = [1.52, 1.04, 0.77, 0.58]

sample_rate = n / full_train_num
print(dataset)
print(n, full_train_num)
print("sample rate", sample_rate)
print("steps", max_token_cnt)

for sigma in sigma_list:
    prv_0 = PoissonSubsampledGaussianMechanism(noise_multiplier=sigma, sampling_probability=sample_rate)
    print("sample rate", sample_rate)

    accountant = PRVAccountant(
        prvs=[prv_0, ],
        max_self_compositions=[1000],
        eps_error=0.01,
        delta_error=1e-10
    )
    eps_low, eps_est, eps_up = accountant.compute_epsilon(delta=1/full_train_num, num_self_compositions=[max_token_cnt])

    print(sigma, eps_low, eps_est, eps_up)

privacy_analysis/subgaussian.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import os
from collections.abc import Sequence, Set

from lmapi.async_tools import limits
from lmapi.auth import OpenAiApiKey
from lmapi.lm import LM, CompletionsSettings
from lmapi.openai import OpenAI

from dp_few_shot_generation.prob_utils import log_normalize


def api_openai_com(model_name: str) -> OpenAI:
    return OpenAI.create(
        "https://api.openai.com/v1/completions",
        # This list is taken from
        # https://github.com/openai/tiktoken/blob/095924e02c85617df6889698d94515f91666c7ea/tiktoken/model.py#L13-L53
        # and modified, currently to accomodate how text-davinci-003 can actually produce <|fim_...|> tokens.
        {
            # chat
            "gpt-4": "cl100k_base",
            "gpt-3.5-turbo": "cl100k_base",
            # text
            "text-davinci-003": "p50k_edit",
            "text-davinci-002": "p50k_base",
            "text-davinci-001": "r50k_base",
            "text-curie-001": "r50k_base",
            "text-babbage-001": "r50k_base",
            "text-ada-001": "r50k_base",
            "davinci": "r50k_base",
            "curie": "r50k_base",
            "babbage": "r50k_base",
            "ada": "r50k_base",
            # code
            "code-davinci-002": "p50k_base",
            "code-davinci-001": "p50k_base",
            "code-cushman-002": "p50k_base",
            "code-cushman-001": "p50k_base",
            "davinci-codex": "p50k_base",
            "cushman-codex": "p50k_base",
            # edit
            "text-davinci-edit-001": "p50k_edit",
            "code-davinci-edit-001": "p50k_edit",
            # embeddings
            "text-embedding-ada-002": "cl100k_base",
            # old embeddings
            "text-similarity-davinci-001": "r50k_base",
            "text-similarity-curie-001": "r50k_base",
            "text-similarity-babbage-001": "r50k_base",
            "text-similarity-ada-001": "r50k_base",
            "text-search-davinci-doc-001": "r50k_base",
            "text-search-curie-doc-001": "r50k_base",
            "text-search-babbage-doc-001": "r50k_base",
            "text-search-ada-doc-001": "r50k_base",
            "code-search-babbage-code-001": "r50k_base",
            "code-search-ada-code-001": "r50k_base",
            # open source
            "gpt2": "gpt2",
        }[model_name],
        OpenAiApiKey(os.environ["OPENAI_API_KEY"]),
        limits.AdaptiveLimiter(),
        {"model": model_name},
    )


MAX_TOP_LOGPROBS = 100
MAX_LOGIT_BIAS = 100
MIN_LOGIT_BIAS = -100


async def next_logprobs(
    self: LM, prompt: str | Sequence[int], top_p=1
) -> dict[int, float]:
    # TODO: Don't hardcode "100" here
    [sampled_tokens] = await self.completions(
        prompt,
        CompletionsSettings(n=1, max_tokens=1, logprobs=100, stop=["<test_for_stop>"]),
    )
    if len(sampled_tokens) == 0:
        if isinstance(prompt, str):
            prompt += "<|endoftext|>"
        else:
            prompt = [*prompt, self.encoding.encode_single_token("<|endoftext|>")]
        [[*_prev_tokens, sampled_token]] = await self.completions(
            prompt,
            CompletionsSettings(
                n=1, max_tokens=0, logprobs=100, echo=True, top_p=top_p
            ),
        )
    else:
        [sampled_token] = sampled_tokens

    return {tlp.token_id: tlp.logprob for tlp in sampled_token.top_choices}


async def normalized_logprobs_for_chosen_tokens(
    self: LM, prompt: Sequence[int], chosen_tokens: Set[int], top_p: float
) -> dict[int, float]:
    """Compute the probability that the prompt will be continued with each of the chosen tokens.

    The returned probability distribution is normalized over just the chosen tokens."""

    assert (
        len(chosen_tokens) <= MAX_TOP_LOGPROBS
    ), f"chosen_tokens must be <= {MAX_TOP_LOGPROBS} in length"

    logit_bias = {token_id: MAX_LOGIT_BIAS for token_id in chosen_tokens}
    [sampled_tokens] = await self.completions(
        prompt,
        CompletionsSettings(
            n=1,
            max_tokens=1,
            logprobs=MAX_TOP_LOGPROBS,
            logit_bias=logit_bias,
            top_p=top_p,
        ),
    )
    if len(sampled_tokens) == 0:
        # Fall back to querying over the set
        chosen_tokens_list = list(chosen_tokens)

        result = await self.completions(
            [[*prompt, token_id] for token_id in chosen_tokens_list],
            CompletionsSettings(n=1, max_tokens=0, logprobs=0, echo=True),
        )
        unnormalized_logprobs = {
            sampled_token.token.token_id: sampled_token.token.logprob
            for [*_prev_tokens, sampled_token] in result
        }
        return log_normalize(unnormalized_logprobs)
    else:
        [sampled_token] = sampled_tokens
        biased_logprobs = {
            tlp.token_id: tlp.logprob for tlp in sampled_token.top_choices
        }
        biased_logprobs_for_tokens = {
            token_id: biased_logprobs.get(token_id, float("-inf"))
            for token_id in chosen_tokens
        }
        return log_normalize(biased_logprobs_for_tokens)

src/dp_few_shot_generation/lm.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

from typing import TypeVar, cast

import numpy as np
import scipy.special

T = TypeVar("T")


def log_normalize(logprobs: dict[T, float]) -> dict[T, float]:
    """Normalize a log probability distribution so that all probabilities sum to 1.

    The input is a sparse distribution represented as a dictionary mapping from token IDs to log probabilities.
    """

    normalizer = cast(float, scipy.special.logsumexp(list(logprobs.values())))
    return {k: v - normalizer for k, v in logprobs.items()}


def log_max_normalize(logprobs: dict[T, float]) -> dict[T, float]:
    """Normalize a log probability distribution so that maxium probabilities is to 1.

    The input is a sparse distribution represented as a dictionary mapping from token IDs to log probabilities.
    """

    normalizer = cast(float, max(list(logprobs.values())))
    return {k: v - normalizer for k, v in logprobs.items()}


def densify(vocab_size: int, logprobs: dict[int, float]) -> np.ndarray:
    """Convert a sparse log-probability distribution into a dense one.

    The dense distribution is represented as a 1D tensor of size `vocab_size`.
    """

    assert len(logprobs) > 0
    result = np.full((vocab_size,), -np.inf)
    for k, v in logprobs.items():
        result[k] = v
    return result


MINIMUM_MISSING_PROB = 1e-6


def remove_logit_bias(
    biased_logprobs: dict[T, float], logit_bias: dict[T, float] | dict[T, int]
) -> dict[T, float]:
    """Undo the effects of logit_bias.

    This function is useful to get log probabilities of arbitrary tokens,
    as a workaround for how OpenAI's API only returns the top K most likely tokens.
    Give those arbitary tokens a large logit_bias (about 70 works well in practice).

    For this function to work properly:
    - The logit_bias should be set so that logprobs for all tokens in the logit_bias
      are provided in biased_logprobs. For example, they should all be large positive numbers.
    - The logit_bias should not be too big, because of limited floating point precision.
      OpenAI's API will not return probabilities smaller than np.log(np.finfo(np.float32).smallest_normal.
    - biased_logprobs should have more elements than logit_bias. Otherwise it's
      not possible to recover the original log probabilities.
    """
    if not (logit_bias.keys() <= biased_logprobs.keys()):
        raise ValueError("logit_bias must be a subset of biased_logprobs")

    missing_key = object()
    missing_prob = 1 - np.exp(scipy.special.logsumexp(list(biased_logprobs.values())))
    missing_logprob = float(
        np.log(missing_prob) if missing_prob > MINIMUM_MISSING_PROB else -np.inf
    )

    debiased_logits: dict[T | object, float] = {
        k: v - logit_bias.get(k, 0) for k, v in biased_logprobs.items()
    }
    debiased_logits[missing_key] = missing_logprob

    normalized = log_normalize(debiased_logits)
    del normalized[missing_key]
    return cast(dict[T, float], normalized)

src/dp_few_shot_generation/prob_utils.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import asyncio
import math
import re
import sys
import time
import traceback
from collections.abc import Iterable, Set
from typing import Annotated, cast

import aiohttp
import more_itertools
import numpy as np
import openai
import scipy.special
import tqdm
import typer
from datasets import DatasetDict, load_dataset
from lmapi.lm import LM, CompletionsSettings
from lmapi.openai import client_session

from dp_few_shot_generation.lm import (
    api_openai_com,
    next_logprobs,
    normalized_logprobs_for_chosen_tokens,
)
from dp_few_shot_generation.prob_utils import densify, log_max_normalize, log_normalize

DEFAULT_NUM_PRIVATE_TRAIN = 20
DEFAULT_NUM_PUBLIC_TRAIN = 0
DEFAULT_NUM_VALID = 4
DEFAULT_NUM_PRIVATE_TRAIN_SPLITS = 10
DEFAULT_NUM_TEST = 1000

labels = ["World", "Sport", "Business", "Technology"]
label_dict = {0: ["World"], 1: ["Sports"], 2: ["Business"], 3: ["Technology"]}


def format_full_datum_for_prompt(labels, datum: dict[str, str]):
    return f'News Type: "{labels[datum["label"]]}"\nText: "{datum["text"] + " END"}"\n'


def format_test_input_for_prompt(labels, test_input: int):
    return f'News Type: "{labels[test_input]}"\nText: "'


def construct_prompt_same(train_examples, test_example):
    prompt = "Classify the news articles into the categories of World, Sports, Business, and Technology.\n\n"
    for train_example in train_examples:
        prompt += "Article: " + train_example["text"] + "\n"
        prompt += "Answer: " + label_dict[train_example["label"]][0] + "\n\n"
    prompt += "Article: " + test_example["text"] + "\n"
    prompt += "Answer:"
    return prompt


def complete(prompt, l, model_name, temp=0, num_log_probs=None, echo=False, n=None):
    # call GPT-3 API until result is provided and then return it
    response = None
    received = False
    while not received:
        try:
            response = openai.Completion.create(
                engine=model_name,
                prompt=prompt,
                max_tokens=l,
                temperature=temp,
                logprobs=num_log_probs,
                echo=echo,
                stop="\n",
                n=n,
            )
            received = True
        except:
            error = sys.exc_info()[0]
            if (
                error == openai.error.InvalidRequestError
            ):  # something is wrong: e.g. prompt too long
                print(f"InvalidRequestError\nPrompt passed in:\n\n{prompt}\n\n")
                assert False

            print("API error:", error)
            time.sleep(1)
    return response


def chunks(lst, n):
    """Yield successive n-sized chunks from lst."""
    for i in range(0, len(lst), n):
        yield lst[i : i + n]


def get_model_response(data, test_examples, openai_model):
    all_raw_answers = []

    prompts = []
    train_examples = data

    for test_example in test_examples:
        prompts.append(construct_prompt_same(train_examples, test_example))

    chunked_prompts = list(chunks(prompts, 20))
    for test_chunk in chunked_prompts:
        response = complete(test_chunk, l=1, model_name=openai_model, num_log_probs=100)

        for answer_id, answer in enumerate(response["choices"]):
            all_raw_answers.append(answer)

    return all_raw_answers


def get_label_probs(all_raw_answers, test_subset):
    """Obtain model's label probability for each of the test examples. The returned prob is NOT normalized"""
    num_classes = len(label_dict)
    approx = False
    assert len(all_raw_answers) == len(test_subset)

    # Fill in the labels that is in the top k prob
    all_label_probs = []
    all_missing_positions = []
    cnt = 0
    for i, ans in enumerate(all_raw_answers):
        try:
            top_logprobs = ans["logprobs"]["top_logprobs"][
                0
            ]  # [0] since we only ask for complete one more token
        except:
            cnt += 1  # cnt for corner case
        label_probs = [0] * len(label_dict.keys())
        for j, label_list in label_dict.items():
            all_found = True
            for label in label_list:  # each possible label correspond to the same class
                label = " " + label  # notice prompt does not have space after 'A:'
                if label in top_logprobs:
                    label_probs[j] += np.exp(top_logprobs[label])
                else:
                    all_found = False
            if not all_found:
                position = (i, j)  # (which test example, which label)
                all_missing_positions.append(position)
        all_label_probs.append(label_probs)
    all_label_probs = np.array(all_label_probs)  # prob not normalized

    return all_label_probs  # NOT NORMALIZED


def eval_accuracy(all_label_probs, test_labels, mode=None, p_cf=None):
    # evaluate the accuracy with and without contextual calibration
    num_classes = all_label_probs.shape[1]
    if p_cf is None:
        # do not calibrate
        W = np.identity(num_classes)
        b = np.zeros([num_classes, 1])
    else:
        # calibrate
        if mode == "diagonal_W":
            W = np.linalg.inv(np.identity(num_classes) * p_cf)
            b = np.zeros([num_classes, 1])
        elif mode == "identity_W":
            W = np.identity(num_classes)
            b = -1 * np.expand_dims(p_cf, axis=-1)
        else:
            assert False

    correctness_list = []
    assert len(all_label_probs) == len(test_labels)
    for label_probs, true_label in zip(all_label_probs, test_labels):
        if np.sum(label_probs) > 0:  # corner case np.sum(label_probs)=0.
            label_probs = label_probs / np.sum(label_probs)  # normalize to 1

        calibrate_label_probs = np.matmul(W, np.expand_dims(label_probs, axis=-1)) + b

        ans_label = np.argmax(calibrate_label_probs)
        if ans_label == true_label:
            correctness_list.append(1)
        else:
            correctness_list.append(0)
    return np.mean(correctness_list)


def get_p_content_free(train_subset, openai_model, content_free_inputs=("N/A",)):
    """Query model with content free input, return its prediction probability for each label"""
    all_p_y = []
    for content_free_input in content_free_inputs:
        prompt = construct_prompt_same(train_subset, content_free_input)
        p_y = [0] * len(label_dict)
        for i, answers in label_dict.items():
            prob = 0
            for a in answers:
                prob += np.exp(
                    complete(
                        prompt + " " + a, 0, openai_model, echo=True, num_log_probs=1
                    )["choices"][0]["logprobs"]["token_logprobs"][-1]
                )
            p_y[i] = prob
        all_p_y.append(p_y)
    p_y = np.mean(np.array(all_p_y), axis=0)
    p_y = p_y / np.sum(p_y)  # normalize
    return p_y


def merge_logprobs_topk_mean(
    private_next_logprobs: list[dict[int, float]],
    public_next_logprobs: dict[int, float],
    n_vocab: int,
    no_public_token: bool,
    normalize_max: bool,
) -> np.ndarray:
    # Compute merged distribution
    # logsumexp - np.log(...): compute mean probability of distribution
    if normalize_max:
        normalize_func = (
            log_max_normalize  # normalize max probability to 1, Exponential mechanism
        )
    else:
        normalize_func = (
            log_normalize  # normalize sum probability to 1, Gaussian mechanism
        )
    if no_public_token:
        merged_next_logprobs = scipy.special.logsumexp(
            np.stack(
                [
                    # Turn into a 1D tensor of size n_vocab
                    densify(
                        n_vocab,
                        # Normalize distribution
                        normalize_func(
                            # Filter to the top 100 most likely next tokens according to the public prompt
                            {k: v for k, v in lps.items()}
                        ),
                    )
                    for lps in private_next_logprobs
                ]
            ),
            axis=0,
        ) - np.log(len(private_next_logprobs))

    else:
        merged_next_logprobs = scipy.special.logsumexp(
            np.stack(
                [
                    # Turn into a 1D tensor of size n_vocab
                    densify(
                        n_vocab,
                        # Normalize distribution
                        normalize_func(
                            # Filter to the top 100 most likely next tokens according to the public prompt
                            {k: v for k, v in lps.items() if k in public_next_logprobs}
                        ),
                    )
                    for lps in private_next_logprobs
                ]
            ),
            axis=0,
        ) - np.log(len(private_next_logprobs))
    merged_next_probs = np.exp(merged_next_logprobs)
    return merged_next_probs


async def generate_with_private_prompts(
    trainset,
    num_private_train,
    num_private_train_splits,
    instruction,
    public_train_prompt: str,
    stop_tokens: Set[int],
    test_input: int,
    lm: LM,
    noise_rng: np.random.RandomState,
    sigma: float,
    labels,
    top_p,
    no_public_token: bool,
    subsample_per_token: bool,
    sample_same_label_prompts: bool,
    gen_seed: int,
    max_tokens: int = 100 - 1,
    normalize_max: bool = False,
) -> list[int]:
    generated_token_ids: list[int] = []

    stringified_test_datum = format_test_input_for_prompt(labels, test_input)
    public_prompt = public_train_prompt + stringified_test_datum
    public_prompt_tokens = lm.encoding.encode(public_prompt)

    assert num_private_train_splits > 0
    if sample_same_label_prompts:
        select_list = []
        for i in range(len(trainset)):
            if trainset[i]["label"] == test_input:
                select_list.append(i)
        train_subset = trainset.select(select_list, keep_in_memory=True)
    else:
        train_subset = trainset.select(range(len(trainset)), keep_in_memory=True)

    if not subsample_per_token:
        private_train_subset = cast(
            Iterable[dict[str, str]],
            train_subset.shuffle(gen_seed, keep_in_memory=True).select(
                range(num_private_train), keep_in_memory=True
            ),
        )
        private_train_splits = [
            list(it)
            for it in more_itertools.distribute(
                num_private_train_splits, private_train_subset
            )
        ]
        private_train_prompts = [
            instruction
            + "\n".join(format_full_datum_for_prompt(labels, datum) for datum in split)
            for split in private_train_splits
        ]
        private_prompts = [
            train_prompt + "\n" + stringified_test_datum
            for train_prompt in private_train_prompts
        ]
        private_prompts_tokens = [
            lm.encoding.encode(prompt) for prompt in private_prompts
        ]

    cnt = 0
    for _ in tqdm.tqdm(range(max_tokens), total=float("inf"), unit=" tokens generated"):
        private_next_logprobs: list[dict[int, float]]
        public_next_logprobs: dict[int, float]
        # Split training dataset
        if subsample_per_token:
            private_train_subset = cast(
                Iterable[dict[str, str]],
                train_subset.shuffle(gen_seed + cnt, keep_in_memory=True).select(
                    range(num_private_train), keep_in_memory=True
                ),
            )
            cnt += 1
            private_train_splits = [
                list(it)
                for it in more_itertools.distribute(
                    num_private_train_splits, private_train_subset
                )
            ]
            # Turn the data into prompts
            private_train_prompts = [
                instruction
                + "\n".join(
                    format_full_datum_for_prompt(labels, datum) for datum in split
                )
                for split in private_train_splits
            ]
            private_prompts = [
                train_prompt + "\n" + stringified_test_datum
                for train_prompt in private_train_prompts
            ]
            private_prompts_tokens = [
                lm.encoding.encode(prompt) for prompt in private_prompts
            ]
        if no_public_token:
            private_next_logprobs = await asyncio.gather(
                *(
                    next_logprobs(lm, prompt + generated_token_ids, top_p=top_p)
                    for prompt in private_prompts_tokens
                )
            )
            merged_next_probs = merge_logprobs_topk_mean(
                private_next_logprobs,
                None,
                lm.encoding.n_vocab,
                no_public_token,
                normalize_max,
            )
            if normalize_max:
                # scale = 1/lambda
                noise = noise_rng.exponential(scale=sigma, size=lm.encoding.n_vocab)
            else:
                noise = noise_rng.normal(0, sigma, size=lm.encoding.n_vocab)
            merged_next_probs += noise
        else:
            public_next_logprobs = await next_logprobs(
                lm, public_prompt_tokens + generated_token_ids, top_p=top_p
            )
            private_next_logprobs = await asyncio.gather(
                *(
                    normalized_logprobs_for_chosen_tokens(
                        lm,
                        prompt + generated_token_ids,
                        public_next_logprobs.keys(),
                        top_p=top_p,
                    )
                    for prompt in private_prompts_tokens
                )
            )
            merged_next_probs = merge_logprobs_topk_mean(
                private_next_logprobs,
                public_next_logprobs,
                lm.encoding.n_vocab,
                no_public_token,
                normalize_max,
            )
            if normalize_max:
                # scale = 1/lambda
                noise = noise_rng.exponential(
                    scale=sigma, size=len(public_next_logprobs)
                )
            else:
                noise = noise_rng.normal(0, sigma, size=len(public_next_logprobs))
            merged_next_probs[list(public_next_logprobs.keys())] += noise

        next_token_id = int(np.argmax(merged_next_probs))

        if next_token_id in stop_tokens:
            break

        generated_token_ids.append(next_token_id)

        del next_token_id
    return generated_token_ids


async def generate_with_public_prompt(
    public_train_prompt: str,
    stop_tokens: Set[str],
    test_input: str,
    lm: LM,
    labels,
    max_tokens: int = 500,
) -> list[int]:
    public_prompt = public_train_prompt + format_test_input_for_prompt(
        labels, test_input
    )
    public_prompt_tokens = lm.encoding.encode(public_prompt)
    public_prompt_tokens = public_prompt

    [completion] = await lm.completions(
        public_prompt_tokens,
        CompletionsSettings(
            temperature=0.0, max_tokens=max_tokens, n=1, stop=list(stop_tokens)
        ),
    )
    generated_tokens = [st.token.token_id for st in completion]
    return generated_tokens


def select_uniform_n_shots_over_labels(data, n_shots):
    select_list = []
    n_shots_per_label = math.ceil(n_shots / len(labels))
    labels_counter = {label[1][0]: n_shots_per_label for label in label_dict.items()}
    n_shots_selected = 0
    for i in range(len(data)):
        label = label_dict[data[i]["label"]][0]
        if labels_counter[label] == 0:
            continue
        else:
            labels_counter[label] -= 1
            select_list.append(i)
            n_shots_selected += 1
        if n_shots_selected == n_shots:
            break
    query_subset = data.select(select_list, keep_in_memory=True)
    return query_subset


def _main(
    sigma: Annotated[float, typer.Option()],  # noise parameters
    openai_model: Annotated[str, typer.Option()] = "babbage",
    print_prompts: Annotated[bool, typer.Option()] = False,
    # num_private_train=MN. MN=0 with num_valid=4 will get epsilon=0 (4-shot) results.
    num_private_train: Annotated[int, typer.Option()] = DEFAULT_NUM_PRIVATE_TRAIN,
    # by default set to 0. set_num_public_train >0 indicates additional public data available.
    set_num_public_train: Annotated[int, typer.Option()] = DEFAULT_NUM_PUBLIC_TRAIN,
    # num_valid=n. n samples to be generated for n-shot ICL
    num_valid: Annotated[int, typer.Option()] = DEFAULT_NUM_VALID,
    # num_private_train_splits=M
    num_private_train_splits: Annotated[
        int, typer.Option()
    ] = DEFAULT_NUM_PRIVATE_TRAIN_SPLITS,
    num_test: Annotated[int, typer.Option()] = DEFAULT_NUM_TEST,
    # no_public_token=True, RVP=False; no_public_token=False, RVP=True
    no_public_token: Annotated[bool, typer.Option()] = False,
    # subsample_per_token=True: at each token generation, subsample a new test set
    subsample_per_token: Annotated[bool, typer.Option()] = False,
    use_dp_prompts: Annotated[bool, typer.Option()] = False,
    # sample_same_label_prompts=True: sample subsets from the sets with same targeted labels.
    sample_same_label_prompts: Annotated[bool, typer.Option()] = False,
    # normalize_max=True, Exponential mechanism; normalize_max=False, Gaussian mechanism
    normalize_max: Annotated[bool, typer.Option()] = False,
    # max_token_per_text=T_max
    max_token_per_text: Annotated[int, typer.Option()] = 100,
    # consistent with default parameters in the documentation https://learn.microsoft.com/en-us/azure/cognitive-services/openai/reference#completions
    top_p: Annotated[float, typer.Option()] = 1,
    # random seed for subsampling in generation
    synth_seed: Annotated[int, typer.Option()] = 0,
    # random seed for n-shot demonstrations sampling in evaluation
    eval_seed: Annotated[int, typer.Option()] = 0,
):
    async def main():
        if (num_private_train == 0) != (num_private_train_splits == 0):
            raise ValueError(
                "Either both or neither of --num-private-train and --num-private-train-splits can be 0"
            )
        command = ["python", sys.argv[0]]
        for x in sys.argv[1:]:
            if x.startswith("--"):
                assert '"' not in x and "'" not in x
                command.append(x)
            else:
                assert "'" not in x
                if re.match("^[a-zA-Z0-9_]+$", x):
                    command.append("%s" % x)
                else:
                    command.append("'%s'" % x)
        command = " ".join(command)
        print(command)

        if no_public_token:
            num_public_train = 0
        else:
            num_public_train = set_num_public_train

        lm = api_openai_com(openai_model)
        noise_rng = np.random.RandomState()

        data = cast(DatasetDict, load_dataset("ag_news"))
        print(labels)

        trainset = data["train"].shuffle(seed=synth_seed, keep_in_memory=True)
        print("trainset length", len(trainset))
        if num_public_train > 0:
            public_train_subset = cast(
                Iterable[dict[str, str]],
                trainset.select(
                    range(
                        len(trainset) - num_public_train,
                        len(trainset),
                        keep_in_memory=True,
                    )
                ),
            )
        else:
            public_train_subset = []

        trainset = trainset.select(
            range(len(trainset) - num_public_train), keep_in_memory=True
        )
        queryset = data["train"].shuffle(seed=eval_seed, keep_in_memory=True)
        query_subset = select_uniform_n_shots_over_labels(queryset, num_valid)

        if use_dp_prompts:
            synthetic_examples = []

            # Turn the data into prompts
            instruction = "Given a label of news type, generate the chosen type of news accordingly.\n\n"

            public_train_prompt = instruction + "\n".join(
                format_full_datum_for_prompt(labels, datum)
                for datum in public_train_subset
            )

            if print_prompts:
                print(public_train_prompt)
                print("=========")

            if normalize_max:
                print("Exponential Mechanism")
                assert num_private_train == 0 or sigma > 0
                if num_private_train > 0:
                    # scale == sigma_calib == 1/lambda. lambda for exponential distribution.
                    sigma_calib = (2 / num_private_train_splits) * (1 / sigma)
            else:
                print("Gaussian Mechanism")
                if num_private_train_splits > 0:
                    sigma_calib = math.sqrt(2) / num_private_train_splits * sigma
                else:
                    sigma_calib = 0
            print(
                f"sigma in command {sigma}. sigma added according to sensitivity {sigma_calib}"
            )

            stop_tokens = {"\n", "<|endoftext|>", " END"}
            stop_tokens_ids = {lm.encoding.encode_single_token(t) for t in stop_tokens}

            client_session.set(aiohttp.ClientSession())

            async with client_session.get():
                for i, test_datum in enumerate(query_subset, 1):
                    print(f"# Example {i}")
                    print(f'News Type: "{labels[test_datum["label"]]}"')
                    print(f'References:\n "{test_datum["text"]}"')

                    np.random.seed(synth_seed + i)
                    gen_seed = np.random.randint(100000)
                    print(f"gen-seed: {gen_seed}")

                    if num_private_train_splits > 0:
                        generated_token_ids = await generate_with_private_prompts(
                            trainset,
                            num_private_train,
                            num_private_train_splits,
                            instruction,
                            public_train_prompt,
                            stop_tokens_ids,
                            test_datum["label"],
                            lm,
                            noise_rng,
                            sigma_calib,
                            labels,
                            top_p,
                            no_public_token,
                            subsample_per_token,
                            sample_same_label_prompts,
                            gen_seed,
                            max_tokens=max_token_per_text
                            - 1,  # need one token length for EOS.
                            normalize_max=normalize_max,
                        )
                    else:
                        generated_token_ids = await generate_with_public_prompt(
                            public_train_prompt,
                            stop_tokens,
                            test_datum["label"],
                            lm,
                            labels,
                            max_tokens=max_token_per_text,
                        )

                    generated = lm.encoding.decode(generated_token_ids).rstrip('"')

                    print(f"Generated: {generated}\n")
                    output_datum = {}
                    output_datum["text"] = generated.strip()
                    output_datum["label"] = test_datum["label"]
                    synthetic_examples.append(output_datum)

        if num_test > 0:
            test_subset = (
                data["test"]
                .shuffle(seed=12345, keep_in_memory=True)
                .select(range(num_test), keep_in_memory=True)
            )
            test_labels = [test_example["label"] for test_example in test_subset]

            content_free_inputs = [{"text": "N/A"}, {"text": ""}, {"text": "[MASK]"}]
            p_cf_wout_DP = get_p_content_free(
                query_subset, openai_model, content_free_inputs=content_free_inputs
            )

            all_raw_answers_wout_DP = get_model_response(
                query_subset, test_subset, openai_model
            )
            all_label_probs_wout_DP = get_label_probs(
                all_raw_answers_wout_DP, test_subset
            )

            acc_original_wout_DP = eval_accuracy(all_label_probs_wout_DP, test_labels)
            acc_calibrated_wout_DP = eval_accuracy(
                all_label_probs_wout_DP,
                test_labels,
                mode="diagonal_W",
                p_cf=p_cf_wout_DP,
            )

            print(f"Accuracy (original) without DP: {acc_original_wout_DP}")
            print(f"Accuracy (calibrated) without DP: {acc_calibrated_wout_DP}")

            if use_dp_prompts:
                p_cf_w_DP = get_p_content_free(
                    synthetic_examples,
                    openai_model,
                    content_free_inputs=content_free_inputs,
                )

                all_raw_answers_w_DP = get_model_response(
                    synthetic_examples, test_subset, openai_model
                )
                all_label_probs_w_DP = get_label_probs(
                    all_raw_answers_w_DP, test_subset
                )

                acc_original_w_DP = eval_accuracy(all_label_probs_w_DP, test_labels)
                acc_calibrated_w_DP = eval_accuracy(
                    all_label_probs_w_DP, test_labels, mode="diagonal_W", p_cf=p_cf_w_DP
                )

                print(f"Accuracy (original) with DP: {acc_original_w_DP}")
                print(f"Accuracy (calibrated) with DP: {acc_calibrated_w_DP}")

    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        traceback.print_exc()
        raise


if __name__ == "__main__":
    typer.run(_main)

src/dp_few_shot_generation/run_exp_agnews.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import asyncio
import math
import re
import sys
import time
import traceback
from collections.abc import Iterable, Set
from typing import Annotated, cast

import aiohttp
import more_itertools
import numpy as np
import openai
import scipy.special
import tqdm
import typer
from datasets import DatasetDict, load_dataset
from lmapi.lm import LM, CompletionsSettings
from lmapi.openai import client_session

from dp_few_shot_generation.lm import (
    api_openai_com,
    next_logprobs,
    normalized_logprobs_for_chosen_tokens,
)
from dp_few_shot_generation.prob_utils import densify, log_max_normalize, log_normalize

DEFAULT_NUM_PRIVATE_TRAIN = 80
DEFAULT_NUM_PUBLIC_TRAIN = 0
DEFAULT_NUM_VALID = 4
DEFAULT_NUM_PRIVATE_TRAIN_SPLITS = 40
DEFAULT_NUM_TEST = 1000

labels = [
    "Company",
    "School",
    "Artist",
    "Ath",
    "Polit",
    "Transportation",
    "Building",
    "Nature",
    "Village",
    "Animal",
    "Plant",
    "Album",
    "Film",
    "Book",
]
label_dict = {
    0: ["Company"],
    1: ["School"],
    2: ["Artist"],
    3: ["Ath"],
    4: ["Polit"],
    5: ["Transportation"],
    6: ["Building"],
    7: ["Nature"],
    8: ["Village"],
    9: ["Animal"],
    10: ["Plant"],
    11: ["Album"],
    12: ["Film"],
    13: ["Book"],
}


def format_full_datum_for_prompt(labels, datum: dict[str, str]):
    return f'Document Type: "{labels[datum["label"]]}"\nText: "{datum["content"] + " END"}"\n'


def format_test_input_for_prompt(labels, test_input: int):
    return f'Document Type: "{labels[test_input]}"\nText: "'


def construct_prompt_same(train_examples, test_example):
    labels_str = ", ".join(labels)
    prompt = (
        f"Classify the documents based on whether they are about a {labels_str}.\n\n"
    )
    for train_example in train_examples:
        prompt += "Article: " + train_example["content"] + "\n"
        prompt += "Answer: " + label_dict[train_example["label"]][0] + "\n\n"
    prompt += "Article: " + test_example["content"] + "\n"
    prompt += "Answer:"
    return prompt


def complete(prompt, l, model_name, temp=0, num_log_probs=None, echo=False, n=None):
    # call GPT-3 API until result is provided and then return it
    response = None
    received = False
    while not received:
        try:
            response = openai.Completion.create(
                engine=model_name,
                prompt=prompt,
                max_tokens=l,
                temperature=temp,
                logprobs=num_log_probs,
                echo=echo,
                stop="\n",
                n=n,
            )
            received = True
        except:
            error = sys.exc_info()[0]
            if (
                error == openai.error.InvalidRequestError
            ):  # something is wrong: e.g. prompt too long
                print(f"InvalidRequestError\nPrompt passed in:\n\n{prompt}\n\n")
                assert False

            print("API error:", error)
            time.sleep(1)
    return response


def chunks(lst, n):
    """Yield successive n-sized chunks from lst."""
    for i in range(0, len(lst), n):
        yield lst[i : i + n]


def get_model_response(data, test_examples, openai_model):
    all_raw_answers = []

    prompts = []
    train_examples = data

    for test_example in test_examples:
        prompts.append(construct_prompt_same(train_examples, test_example))

    chunked_prompts = list(chunks(prompts, 20))
    for test_chunk in chunked_prompts:
        response = complete(test_chunk, l=1, model_name=openai_model, num_log_probs=100)

        for answer_id, answer in enumerate(response["choices"]):
            all_raw_answers.append(answer)

    return all_raw_answers


def get_label_probs(all_raw_answers, test_subset):
    """Obtain model's label probability for each of the test examples. The returned prob is NOT normalized"""
    num_classes = len(label_dict)
    approx = False
    assert len(all_raw_answers) == len(test_subset)

    # Fill in the labels that is in the top k prob
    all_label_probs = []
    all_missing_positions = []
    cnt = 0
    for i, ans in enumerate(all_raw_answers):
        try:
            top_logprobs = ans["logprobs"]["top_logprobs"][
                0
            ]  # [0] since we only ask for complete one more token
        except:
            cnt += 1  # cnt for corner case
        label_probs = [0] * len(label_dict.keys())
        for j, label_list in label_dict.items():
            all_found = True
            for label in label_list:  # each possible label correspond to the same class
                label = " " + label  # notice prompt does not have space after 'A:'
                if label in top_logprobs:
                    label_probs[j] += np.exp(top_logprobs[label])
                else:
                    all_found = False
            if not all_found:
                position = (i, j)  # (which test example, which label)
                all_missing_positions.append(position)
        all_label_probs.append(label_probs)
    all_label_probs = np.array(all_label_probs)  # prob not normalized

    return all_label_probs  # NOT NORMALIZED


def eval_accuracy(all_label_probs, test_labels, mode=None, p_cf=None):
    # evaluate the accuracy with and without contextual calibration
    num_classes = all_label_probs.shape[1]
    if p_cf is None:
        # do not calibrate
        W = np.identity(num_classes)
        b = np.zeros([num_classes, 1])
    else:
        # calibrate
        if mode == "diagonal_W":
            W = np.linalg.inv(np.identity(num_classes) * p_cf)
            b = np.zeros([num_classes, 1])
        elif mode == "identity_W":
            W = np.identity(num_classes)
            b = -1 * np.expand_dims(p_cf, axis=-1)
        else:
            assert False

    correctness_list = []
    assert len(all_label_probs) == len(test_labels)
    for label_probs, true_label in zip(all_label_probs, test_labels):
        if np.sum(label_probs) > 0:  # corner case np.sum(label_probs)=0.
            label_probs = label_probs / np.sum(label_probs)  # normalize to 1

        calibrate_label_probs = np.matmul(W, np.expand_dims(label_probs, axis=-1)) + b

        ans_label = np.argmax(calibrate_label_probs)
        if ans_label == true_label:
            correctness_list.append(1)
        else:
            correctness_list.append(0)
    return np.mean(correctness_list)


def get_p_content_free(train_subset, openai_model, content_free_inputs=("N/A",)):
    """Query model with content free input, return its prediction probability for each label"""
    all_p_y = []
    for content_free_input in content_free_inputs:
        prompt = construct_prompt_same(train_subset, content_free_input)
        p_y = [0] * len(label_dict)
        for i, answers in label_dict.items():
            prob = 0
            for a in answers:
                prob += np.exp(
                    complete(
                        prompt + " " + a, 0, openai_model, echo=True, num_log_probs=1
                    )["choices"][0]["logprobs"]["token_logprobs"][-1]
                )
            p_y[i] = prob
        all_p_y.append(p_y)
    p_y = np.mean(np.array(all_p_y), axis=0)
    p_y = p_y / np.sum(p_y)  # normalize
    return p_y


def merge_logprobs_topk_mean(
    private_next_logprobs: list[dict[int, float]],
    public_next_logprobs: dict[int, float],
    n_vocab: int,
    no_public_token: bool,
    normalize_max: bool,
) -> np.ndarray:
    # Compute merged distribution
    # logsumexp - np.log(...): compute mean probability of distribution
    if normalize_max:
        normalize_func = (
            log_max_normalize  # normalize max probability to 1, Exponential mechanism
        )
    else:
        normalize_func = (
            log_normalize  # normalize sum probability to 1, Gaussian mechanism
        )
    if no_public_token:
        merged_next_logprobs = scipy.special.logsumexp(
            np.stack(
                [
                    # Turn into a 1D tensor of size n_vocab
                    densify(
                        n_vocab,
                        # Normalize distribution
                        normalize_func(
                            # Filter to the top 100 most likely next tokens according to the public prompt
                            {k: v for k, v in lps.items()}
                        ),
                    )
                    for lps in private_next_logprobs
                ]
            ),
            axis=0,
        ) - np.log(len(private_next_logprobs))

    else:
        merged_next_logprobs = scipy.special.logsumexp(
            np.stack(
                [
                    # Turn into a 1D tensor of size n_vocab
                    densify(
                        n_vocab,
                        # Normalize distribution
                        normalize_func(
                            # Filter to the top 100 most likely next tokens according to the public prompt
                            {k: v for k, v in lps.items() if k in public_next_logprobs}
                        ),
                    )
                    for lps in private_next_logprobs
                ]
            ),
            axis=0,
        ) - np.log(len(private_next_logprobs))
    merged_next_probs = np.exp(merged_next_logprobs)
    return merged_next_probs


async def generate_with_private_prompts(
    trainset,
    num_private_train,
    num_private_train_splits,
    instruction,
    public_train_prompt: str,
    stop_tokens: Set[int],
    test_input: int,
    lm: LM,
    noise_rng: np.random.RandomState,
    sigma: float,
    labels,
    top_p,
    no_public_token: bool,
    subsample_per_token: bool,
    sample_same_label_prompts: bool,
    gen_seed: int,
    max_tokens: int,
    normalize_max: bool = False,
) -> list[int]:
    generated_token_ids: list[int] = []

    stringified_test_datum = format_test_input_for_prompt(labels, test_input)
    public_prompt = public_train_prompt + stringified_test_datum
    public_prompt_tokens = lm.encoding.encode(public_prompt)

    assert num_private_train_splits > 0
    if sample_same_label_prompts:
        select_list = []
        for i in range(len(trainset)):
            if trainset[i]["label"] == test_input:
                select_list.append(i)
        train_subset = trainset.select(select_list, keep_in_memory=True)
    else:
        train_subset = trainset.select(range(len(trainset)), keep_in_memory=True)

    if not subsample_per_token:
        private_train_subset = cast(
            Iterable[dict[str, str]],
            train_subset.shuffle(gen_seed, keep_in_memory=True).select(
                range(num_private_train), keep_in_memory=True
            ),
        )
        private_train_splits = [
            list(it)
            for it in more_itertools.distribute(
                num_private_train_splits, private_train_subset
            )
        ]
        private_train_prompts = [
            instruction
            + "\n".join(format_full_datum_for_prompt(labels, datum) for datum in split)
            for split in private_train_splits
        ]
        private_prompts = [
            train_prompt + "\n" + stringified_test_datum
            for train_prompt in private_train_prompts
        ]
        private_prompts_tokens = [
            lm.encoding.encode(prompt) for prompt in private_prompts
        ]

    cnt = 0
    for _ in tqdm.tqdm(range(max_tokens), total=float("inf"), unit=" tokens generated"):
        private_next_logprobs: list[dict[int, float]]
        public_next_logprobs: dict[int, float]
        # Split training dataset
        if subsample_per_token:
            private_train_subset = cast(
                Iterable[dict[str, str]],
                train_subset.shuffle(gen_seed + cnt, keep_in_memory=True).select(
                    range(num_private_train), keep_in_memory=True
                ),
            )
            cnt += 1
            private_train_splits = [
                list(it)
                for it in more_itertools.distribute(
                    num_private_train_splits, private_train_subset
                )
            ]
            # Turn the data into prompts
            private_train_prompts = [
                instruction
                + "\n".join(
                    format_full_datum_for_prompt(labels, datum) for datum in split
                )
                for split in private_train_splits
            ]
            private_prompts = [
                train_prompt + "\n" + stringified_test_datum
                for train_prompt in private_train_prompts
            ]
            private_prompts_tokens = [
                lm.encoding.encode(prompt) for prompt in private_prompts
            ]
        if no_public_token:
            private_next_logprobs = await asyncio.gather(
                *(
                    next_logprobs(lm, prompt + generated_token_ids, top_p=top_p)
                    for prompt in private_prompts_tokens
                )
            )
            merged_next_probs = merge_logprobs_topk_mean(
                private_next_logprobs,
                None,
                lm.encoding.n_vocab,
                no_public_token,
                normalize_max,
            )
            if normalize_max:
                # scale = 1/lambda
                noise = noise_rng.exponential(scale=sigma, size=lm.encoding.n_vocab)
            else:
                noise = noise_rng.normal(0, sigma, size=lm.encoding.n_vocab)
            merged_next_probs += noise
        else:
            public_next_logprobs = await next_logprobs(
                lm, public_prompt_tokens + generated_token_ids, top_p=top_p
            )
            private_next_logprobs = await asyncio.gather(
                *(
                    normalized_logprobs_for_chosen_tokens(
                        lm,
                        prompt + generated_token_ids,
                        public_next_logprobs.keys(),
                        top_p=top_p,
                    )
                    for prompt in private_prompts_tokens
                )
            )
            merged_next_probs = merge_logprobs_topk_mean(
                private_next_logprobs,
                public_next_logprobs,
                lm.encoding.n_vocab,
                no_public_token,
                normalize_max,
            )
            if normalize_max:
                # scale = 1/lambda
                noise = noise_rng.exponential(
                    scale=sigma, size=len(public_next_logprobs)
                )
            else:
                noise = noise_rng.normal(0, sigma, size=len(public_next_logprobs))
            merged_next_probs[list(public_next_logprobs.keys())] += noise

        next_token_id = int(np.argmax(merged_next_probs))

        if next_token_id in stop_tokens:
            break

        generated_token_ids.append(next_token_id)

        del next_token_id
    return generated_token_ids


async def generate_with_public_prompt(
    public_train_prompt: str,
    stop_tokens: Set[str],
    test_input: str,
    lm: LM,
    labels,
    max_tokens: int = 500,
) -> list[int]:
    public_prompt = public_train_prompt + format_test_input_for_prompt(
        labels, test_input
    )
    public_prompt_tokens = lm.encoding.encode(public_prompt)
    public_prompt_tokens = public_prompt

    [completion] = await lm.completions(
        public_prompt_tokens,
        CompletionsSettings(
            temperature=0.0, max_tokens=max_tokens, n=1, stop=list(stop_tokens)
        ),
    )
    generated_tokens = [st.token.token_id for st in completion]
    return generated_tokens


def select_uniform_n_shots_over_labels(data, n_shots):
    select_list = []
    n_shots_per_label = math.ceil(n_shots / len(labels))
    labels_counter = {label[1][0]: n_shots_per_label for label in label_dict.items()}
    n_shots_selected = 0
    for i in range(len(data)):
        label = label_dict[data[i]["label"]][0]
        if labels_counter[label] == 0:
            continue
        else:
            labels_counter[label] -= 1
            select_list.append(i)
            n_shots_selected += 1
        if n_shots_selected == n_shots:
            break
    query_subset = data.select(select_list, keep_in_memory=True)
    return query_subset


def _main(
    sigma: Annotated[float, typer.Option()],  # noise parameters
    openai_model: Annotated[str, typer.Option()] = "babbage",
    print_prompts: Annotated[bool, typer.Option()] = False,
    # num_private_train=MN. MN=0 with num_valid=4 will get epsilon=0 (4-shot) results.
    num_private_train: Annotated[int, typer.Option()] = DEFAULT_NUM_PRIVATE_TRAIN,
    # by default set to 0. set_num_public_train >0 indicates additional public data available.
    set_num_public_train: Annotated[int, typer.Option()] = DEFAULT_NUM_PUBLIC_TRAIN,
    # num_valid=n. n samples to be generated for n-shot ICL
    num_valid: Annotated[int, typer.Option()] = DEFAULT_NUM_VALID,
    # num_private_train_splits=M
    num_private_train_splits: Annotated[
        int, typer.Option()
    ] = DEFAULT_NUM_PRIVATE_TRAIN_SPLITS,
    num_test: Annotated[int, typer.Option()] = DEFAULT_NUM_TEST,
    # no_public_token=True, RVP=False; no_public_token=False, RVP=True
    no_public_token: Annotated[bool, typer.Option()] = False,
    # subsample_per_token=True: at each token generation, subsample a new test set
    subsample_per_token: Annotated[bool, typer.Option()] = False,
    use_dp_prompts: Annotated[bool, typer.Option()] = False,
    # sample_same_label_prompts=True: sample subsets from the sets with same targeted labels.
    sample_same_label_prompts: Annotated[bool, typer.Option()] = False,
    # normalize_max=True, Exponential mechanism; normalize_max=False, Gaussian mechanism
    normalize_max: Annotated[bool, typer.Option()] = False,
    # max_token_per_text=T_max
    max_token_per_text: Annotated[int, typer.Option()] = 100,
    # consistent with default parameters in the documentation https://learn.microsoft.com/en-us/azure/cognitive-services/openai/reference#completions
    top_p: Annotated[float, typer.Option()] = 1,
    # random seed for subsampling in generation
    synth_seed: Annotated[int, typer.Option()] = 0,
    # random seed for n-shot demonstrations sampling in evaluation
    eval_seed: Annotated[int, typer.Option()] = 0,
):
    async def main():
        if (num_private_train == 0) != (num_private_train_splits == 0):
            raise ValueError(
                "Either both or neither of --num-private-train and --num-private-train-splits can be 0"
            )
        command = ["python", sys.argv[0]]
        for x in sys.argv[1:]:
            if x.startswith("--"):
                assert '"' not in x and "'" not in x
                command.append(x)
            else:
                assert "'" not in x
                if re.match("^[a-zA-Z0-9_]+$", x):
                    command.append("%s" % x)
                else:
                    command.append("'%s'" % x)
        command = " ".join(command)
        print(command)

        if no_public_token:
            num_public_train = 0
        else:
            num_public_train = set_num_public_train

        lm = api_openai_com(openai_model)
        noise_rng = np.random.RandomState()

        data = cast(DatasetDict, load_dataset("dbpedia_14"))
        print(labels)

        trainset = data["train"].shuffle(seed=synth_seed, keep_in_memory=True)
        print("trainset length", len(trainset))
        if num_public_train > 0:
            public_train_subset = cast(
                Iterable[dict[str, str]],
                trainset.select(
                    range(
                        len(trainset) - num_public_train,
                        len(trainset),
                        keep_in_memory=True,
                    )
                ),
            )
        else:
            public_train_subset = []

        trainset = trainset.select(
            range(len(trainset) - num_public_train), keep_in_memory=True
        )
        queryset = data["train"].shuffle(seed=eval_seed, keep_in_memory=True)
        query_subset = select_uniform_n_shots_over_labels(queryset, num_valid)

        if use_dp_prompts:
            synthetic_examples = []

            # Turn the data into prompts
            instruction = "Given a label of document type, generate the chosen type of document accordingly.\n\n"

            public_train_prompt = instruction + "\n".join(
                format_full_datum_for_prompt(labels, datum)
                for datum in public_train_subset
            )

            if print_prompts:
                print(public_train_prompt)
                print("=========")

            if normalize_max:
                print("Exponential Mechanism")
                assert num_private_train == 0 or sigma > 0
                if num_private_train > 0:
                    # scale == sigma_calib == 1/lambda. lambda for exponential distribution.
                    sigma_calib = (2 / num_private_train_splits) * (1 / sigma)
            else:
                print("Gaussian Mechanism")
                if num_private_train_splits > 0:
                    sigma_calib = math.sqrt(2) / num_private_train_splits * sigma
                else:
                    sigma_calib = 0
            print(
                f"sigma in command {sigma}. sigma added according to sensitivity {sigma_calib}"
            )

            stop_tokens = {"\n", "<|endoftext|>", " END"}
            stop_tokens_ids = {lm.encoding.encode_single_token(t) for t in stop_tokens}

            client_session.set(aiohttp.ClientSession())

            async with client_session.get():
                for i, test_datum in enumerate(query_subset, 1):
                    print(f"# Example {i}")
                    print(f'Document Type: "{labels[test_datum["label"]]}"')
                    print(f'References:\n "{test_datum["content"]}"')

                    np.random.seed(synth_seed + i)
                    gen_seed = np.random.randint(100000)
                    print(f"gen-seed: {gen_seed}")

                    if num_private_train_splits > 0:
                        generated_token_ids = await generate_with_private_prompts(
                            trainset,
                            num_private_train,
                            num_private_train_splits,
                            instruction,
                            public_train_prompt,
                            stop_tokens_ids,
                            test_datum["label"],
                            lm,
                            noise_rng,
                            sigma_calib,
                            labels,
                            top_p,
                            no_public_token,
                            subsample_per_token,
                            sample_same_label_prompts,
                            gen_seed,
                            max_tokens=max_token_per_text
                            - 1,  # need one token length for EOS.
                            normalize_max=normalize_max,
                        )
                    else:
                        generated_token_ids = await generate_with_public_prompt(
                            public_train_prompt,
                            stop_tokens,
                            test_datum["label"],
                            lm,
                            labels,
                            max_tokens=max_token_per_text,
                        )

                    generated = lm.encoding.decode(generated_token_ids).rstrip('"')

                    print(f"Generated: {generated}\n")
                    output_datum = {}
                    output_datum["content"] = generated.strip()
                    output_datum["label"] = test_datum["label"]
                    synthetic_examples.append(output_datum)

        if num_test > 0:
            test_subset = (
                data["test"]
                .shuffle(seed=12345, keep_in_memory=True)
                .select(range(num_test), keep_in_memory=True)
            )
            test_labels = [test_example["label"] for test_example in test_subset]

            content_free_inputs = [
                {"content": "N/A"},
                {"content": ""},
                {"content": "[MASK]"},
            ]
            p_cf_wout_DP = get_p_content_free(
                query_subset, openai_model, content_free_inputs=content_free_inputs
            )

            all_raw_answers_wout_DP = get_model_response(
                query_subset, test_subset, openai_model
            )
            all_label_probs_wout_DP = get_label_probs(
                all_raw_answers_wout_DP, test_subset
            )

            acc_original_wout_DP = eval_accuracy(all_label_probs_wout_DP, test_labels)
            acc_calibrated_wout_DP = eval_accuracy(
                all_label_probs_wout_DP,
                test_labels,
                mode="diagonal_W",
                p_cf=p_cf_wout_DP,
            )

            print(f"Accuracy (original) without DP: {acc_original_wout_DP}")
            print(f"Accuracy (calibrated) without DP: {acc_calibrated_wout_DP}")

            if use_dp_prompts:
                p_cf_w_DP = get_p_content_free(
                    synthetic_examples,
                    openai_model,
                    content_free_inputs=content_free_inputs,
                )
                all_raw_answers_w_DP = get_model_response(
                    synthetic_examples, test_subset, openai_model
                )

                all_label_probs_w_DP = get_label_probs(
                    all_raw_answers_w_DP, test_subset
                )

                acc_original_w_DP = eval_accuracy(all_label_probs_w_DP, test_labels)
                acc_calibrated_w_DP = eval_accuracy(
                    all_label_probs_w_DP, test_labels, mode="diagonal_W", p_cf=p_cf_w_DP
                )

                print(f"Accuracy (original) with DP: {acc_original_w_DP}")
                print(f"Accuracy (calibrated) with DP: {acc_calibrated_w_DP}")

    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        traceback.print_exc()
        raise


if __name__ == "__main__":
    typer.run(_main)

src/dp_few_shot_generation/run_exp_dbpedia.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import asyncio
import math
import re
import sys
import traceback
from collections.abc import Iterable, Set
from typing import Annotated, cast

import aiohttp
import more_itertools
import numpy as np
import openai
import scipy.special
import tqdm
import typer
from datasets import DatasetDict, load_dataset
from lmapi.lm import LM, CompletionsSettings
from lmapi.openai import client_session

from dp_few_shot_generation.lm import (
    api_openai_com,
    next_logprobs,
    normalized_logprobs_for_chosen_tokens,
)
from dp_few_shot_generation.prob_utils import densify, log_max_normalize, log_normalize

DEFAULT_NUM_PRIVATE_TRAIN = 80
DEFAULT_NUM_PUBLIC_TRAIN = 0
DEFAULT_NUM_VALID = 4
DEFAULT_NUM_PRIVATE_TRAIN_SPLITS = 20
DEFAULT_NUM_TEST = -1


def format_full_datum_for_prompt(field_name, datum: dict[str, str]):
    return (
        f'{field_name}: "{datum["label"]}"\nSentence: "{datum["content"] + " END"}"\n'
    )


def format_test_input_for_prompt(field_name, test_input: str):
    return f'{field_name}: "{test_input}"\nSentence: "'


def construct_prompt_same(train_examples, test_example, field_name):
    prompt = f""  # prompt strucrture follows: https://github.com/tonyzhaozh/few-shot-learning/blob/main/data_utils.py#L427-L429
    for train_example in train_examples:
        prompt += "Sentence: " + train_example["content"] + "\n"
        prompt += f"{field_name}: " + train_example["label"] + "\n\n"
    prompt += "Sentence: " + test_example["content"] + "\n"
    prompt += f"{field_name}:"
    return prompt


def complete(prompt, l, model_name, temp=0, num_log_probs=None, echo=False, n=None):
    # call GPT-3 API until result is provided and then return it
    response = None
    received = False
    while not received:
        try:
            response = openai.Completion.create(
                engine=model_name,
                prompt=prompt,
                max_tokens=l,
                temperature=temp,
                logprobs=num_log_probs,
                echo=echo,
                stop="\n",
                n=n,
            )
            received = True
        except:
            error = sys.exc_info()[0]
            if (
                error == openai.error.InvalidRequestError
            ):  # something is wrong: e.g. prompt too long
                print(f"InvalidRequestError\nPrompt passed in:\n\n{prompt}\n\n")
                assert False

            print("API error:", error)
            time.sleep(1)
    return response


def chunks(lst, n):
    """Yield successive n-sized chunks from lst."""
    for i in range(0, len(lst), n):
        yield lst[i : i + n]


def get_model_response(
    data,
    test_examples,
    openai_model,
    field_name,
    max_token_to_fill=5,
    additional_tokens=None,
):
    all_raw_answers = []

    prompts = []
    train_examples = data

    for test_example in test_examples:
        prompts.append(construct_prompt_same(train_examples, test_example, field_name))

    if additional_tokens is not None:
        assert len(additional_tokens) == len(prompts)
        for i in range(len(prompts)):
            prompts[i] += additional_tokens[i]

    chunked_prompts = list(chunks(prompts, 20))
    for test_chunk in chunked_prompts:
        response = complete(
            test_chunk, l=max_token_to_fill, model_name=openai_model, num_log_probs=100
        )

        for answer_id, answer in enumerate(response["choices"]):
            all_raw_answers.append(answer)

    return all_raw_answers


def em_accuracy_helper(prediction, label):
    correctness_list = []
    for pred, l in zip(prediction, label):
        pred = pred.split("\n")[0]
        if pred == l:
            correctness_list.append(1)
        else:
            correctness_list.append(0)
    return np.mean(correctness_list)


def merge_logprobs_topk_mean(
    private_next_logprobs: list[dict[int, float]],
    public_next_logprobs: dict[int, float],
    n_vocab: int,
    no_public_token: bool,
    normalize_max: bool,
) -> np.ndarray:
    # Compute merged distribution
    # logsumexp - np.log(...): compute mean probability of distribution
    if normalize_max:
        normalize_func = (
            log_max_normalize  # normalize max probability to 1, Exponential mechanism
        )
    else:
        normalize_func = (
            log_normalize  # normalize sum probability to 1, Gaussian mechanism
        )
    if no_public_token:
        merged_next_logprobs = scipy.special.logsumexp(
            np.stack(
                [
                    # Turn into a 1D tensor of size n_vocab
                    densify(
                        n_vocab,
                        # Normalize distribution
                        normalize_func(
                            # Filter to the top 100 most likely next tokens according to the public prompt
                            {k: v for k, v in lps.items()}
                        ),
                    )
                    for lps in private_next_logprobs
                ]
            ),
            axis=0,
        ) - np.log(len(private_next_logprobs))

    else:
        merged_next_logprobs = scipy.special.logsumexp(
            np.stack(
                [
                    # Turn into a 1D tensor of size n_vocab
                    densify(
                        n_vocab,
                        # Normalize distribution
                        normalize_func(
                            # Filter to the top 100 most likely next tokens according to the public prompt
                            {k: v for k, v in lps.items() if k in public_next_logprobs}
                        ),
                    )
                    for lps in private_next_logprobs
                ]
            ),
            axis=0,
        ) - np.log(len(private_next_logprobs))
    merged_next_probs = np.exp(merged_next_logprobs)
    return merged_next_probs


async def generate_with_private_prompts(
    trainset,
    num_private_train,
    num_private_train_splits,
    instruction,
    public_train_prompt: str,
    stop_tokens: Set[int],
    test_input: str,
    lm: LM,
    noise_rng: np.random.RandomState,
    sigma: float,
    field_name: str,
    top_p,
    no_public_token: bool,
    subsample_per_token: bool,
    gen_seed: int,
    max_tokens: int,
    normalize_max: bool = False,
) -> list[int]:
    generated_token_ids: list[int] = []

    stringified_test_datum = format_test_input_for_prompt(field_name, test_input)
    public_prompt = public_train_prompt + stringified_test_datum
    public_prompt_tokens = lm.encoding.encode(public_prompt)

    assert num_private_train_splits > 0

    train_subset = trainset.select(range(len(trainset)), keep_in_memory=True)

    if not subsample_per_token:
        private_train_subset = cast(
            Iterable[dict[str, str]],
            train_subset.shuffle(gen_seed, keep_in_memory=True).select(
                range(num_private_train), keep_in_memory=True
            ),
        )
        private_train_splits = [
            list(it)
            for it in more_itertools.distribute(
                num_private_train_splits, private_train_subset
            )
        ]
        private_train_prompts = [
            instruction
            + "\n".join(
                format_full_datum_for_prompt(field_name, datum) for datum in split
            )
            for split in private_train_splits
        ]
        private_prompts = [
            train_prompt + "\n" + stringified_test_datum
            for train_prompt in private_train_prompts
        ]
        private_prompts_tokens = [
            lm.encoding.encode(prompt) for prompt in private_prompts
        ]

    cnt = 0
    for _ in tqdm.tqdm(range(max_tokens), total=float("inf"), unit=" tokens generated"):
        private_next_logprobs: list[dict[int, float]]
        public_next_logprobs: dict[int, float]
        # Split training dataset
        if subsample_per_token:
            private_train_subset = cast(
                Iterable[dict[str, str]],
                train_subset.shuffle(gen_seed + cnt, keep_in_memory=True).select(
                    range(num_private_train), keep_in_memory=True
                ),
            )
            cnt += 1
            private_train_splits = [
                list(it)
                for it in more_itertools.distribute(
                    num_private_train_splits, private_train_subset
                )
            ]
            # Turn the data into prompts
            private_train_prompts = [
                instruction
                + "\n".join(
                    format_full_datum_for_prompt(field_name, datum) for datum in split
                )
                for split in private_train_splits
            ]
            private_prompts = [
                train_prompt + "\n" + stringified_test_datum
                for train_prompt in private_train_prompts
            ]
            private_prompts_tokens = [
                lm.encoding.encode(prompt) for prompt in private_prompts
            ]
        if no_public_token:
            private_next_logprobs = await asyncio.gather(
                *(
                    next_logprobs(lm, prompt + generated_token_ids, top_p=top_p)
                    for prompt in private_prompts_tokens
                )
            )
            merged_next_probs = merge_logprobs_topk_mean(
                private_next_logprobs,
                None,
                lm.encoding.n_vocab,
                no_public_token,
                normalize_max,
            )

            if normalize_max:
                # scale = 1/lambda
                noise = noise_rng.exponential(scale=sigma, size=lm.encoding.n_vocab)
            else:
                noise = noise_rng.normal(0, sigma, size=lm.encoding.n_vocab)
            merged_next_probs += noise
        else:
            public_next_logprobs = await next_logprobs(
                lm, public_prompt_tokens + generated_token_ids, top_p=top_p
            )
            private_next_logprobs = await asyncio.gather(
                *(
                    normalized_logprobs_for_chosen_tokens(
                        lm,
                        prompt + generated_token_ids,
                        public_next_logprobs.keys(),
                        top_p=top_p,
                    )
                    for prompt in private_prompts_tokens
                )
            )
            merged_next_probs = merge_logprobs_topk_mean(
                private_next_logprobs,
                public_next_logprobs,
                lm.encoding.n_vocab,
                no_public_token,
                normalize_max,
            )
            if normalize_max:
                # scale = 1/lambda
                noise = noise_rng.exponential(
                    scale=sigma, size=len(public_next_logprobs)
                )
            else:
                noise = noise_rng.normal(0, sigma, size=len(public_next_logprobs))
            merged_next_probs[list(public_next_logprobs.keys())] += noise

        next_token_id = int(np.argmax(merged_next_probs))

        if next_token_id in stop_tokens:
            break

        generated_token_ids.append(next_token_id)

        del next_token_id
    return generated_token_ids


async def generate_with_public_prompt(
    public_train_prompt: str,
    stop_tokens: Set[str],
    test_input: str,
    lm: LM,
    field_name,
    max_tokens: int = 500,
) -> list[int]:
    public_prompt = public_train_prompt + format_test_input_for_prompt(
        field_name, test_input
    )
    public_prompt_tokens = lm.encoding.encode(public_prompt)
    public_prompt_tokens = public_prompt

    [completion] = await lm.completions(
        public_prompt_tokens,
        CompletionsSettings(
            temperature=0.0, max_tokens=max_tokens, n=1, stop=list(stop_tokens)
        ),
    )
    generated_tokens = [st.token.token_id for st in completion]
    return generated_tokens


def _main(
    sigma: Annotated[float, typer.Option()],  # noise parameters
    openai_model: Annotated[str, typer.Option()] = "babbage",
    print_prompts: Annotated[bool, typer.Option()] = False,
    # num_private_train=MN. MN=0 with num_valid=4 will get epsilon=0 (4-shot) results.
    num_private_train: Annotated[int, typer.Option()] = DEFAULT_NUM_PRIVATE_TRAIN,
    # by default set to 0. set_num_public_train >0 indicates additional public data available.
    set_num_public_train: Annotated[int, typer.Option()] = DEFAULT_NUM_PUBLIC_TRAIN,
    # num_valid=n. n samples to be generated for n-shot ICL
    num_valid: Annotated[int, typer.Option()] = DEFAULT_NUM_VALID,
    # num_private_train_splits=M
    num_private_train_splits: Annotated[
        int, typer.Option()
    ] = DEFAULT_NUM_PRIVATE_TRAIN_SPLITS,
    num_test: Annotated[int, typer.Option()] = DEFAULT_NUM_TEST,
    # no_public_token=True, RVP=False; no_public_token=False, RVP=True
    no_public_token: Annotated[bool, typer.Option()] = False,
    # subsample_per_token=True: at each token generation, subsample a new test set
    subsample_per_token: Annotated[bool, typer.Option()] = False,
    use_dp_prompts: Annotated[bool, typer.Option()] = False,
    # normalize_max=True, Exponential mechanism; normalize_max=False, Gaussian mechanism
    normalize_max: Annotated[bool, typer.Option()] = False,
    # max_token_per_text=T_max
    max_token_per_text: Annotated[int, typer.Option()] = 20,
    # consistent with default parameters in the documentation https://learn.microsoft.com/en-us/azure/cognitive-services/openai/reference#completions
    top_p: Annotated[float, typer.Option()] = 1,
    # random seed for subsampling in generation
    synth_seed: Annotated[int, typer.Option()] = 0,
    # random seed for n-shot demonstrations sampling in evaluation
    eval_seed: Annotated[int, typer.Option()] = 0,
    # choice bewteen ["Genre", "Director"]
    field_name: Annotated[str, typer.Option()] = "Genre",
    data_path: Annotated[str, typer.Option()] = "./../../data/movie",
):
    async def main():
        if (num_private_train == 0) != (num_private_train_splits == 0):
            raise ValueError(
                "Either both or neither of --num-private-train and --num-private-train-splits can be 0"
            )
        assert field_name in [
            "Director",
            "Genre",
        ]  # field_name from movie dataset include "Actor", "Award", "Character_Name", "Director", "Genre", "Opinion", "Origin", "Plot", "Quote", "Relationship", "Soundtrack", "Year"]
        command = ["python", sys.argv[0]]
        for x in sys.argv[1:]:
            if x.startswith("--"):
                assert '"' not in x and "'" not in x
                command.append(x)
            else:
                assert "'" not in x
                if re.match("^[a-zA-Z0-9_]+$", x):
                    command.append("%s" % x)
                else:
                    command.append("'%s'" % x)
        command = " ".join(command)
        print(command)

        if no_public_token:
            num_public_train = 0
        else:
            num_public_train = set_num_public_train

        lm = api_openai_com(openai_model)
        noise_rng = np.random.RandomState()

        data_files = {"train": "train.csv", "test": "test.csv"}
        data = cast(
            DatasetDict,
            load_dataset(f"{data_path}/{field_name}", data_files=data_files),
        )

        trainset = data["train"].shuffle(seed=synth_seed, keep_in_memory=True)
        print("trainset length", len(trainset))
        if num_public_train > 0:
            public_train_subset = cast(
                Iterable[dict[str, str]],
                trainset.select(
                    range(
                        len(trainset) - num_public_train,
                        len(trainset),
                        keep_in_memory=True,
                    )
                ),
            )
        else:
            public_train_subset = []

        trainset = trainset.select(
            range(len(trainset) - num_public_train), keep_in_memory=True
        )
        query_subset = (
            data["train"]
            .shuffle(seed=eval_seed, keep_in_memory=True)
            .select(range(num_valid), keep_in_memory=True)
        )

        if use_dp_prompts:
            synthetic_examples = []

            # Turn the data into prompts
            instruction = f"Given a propety of {field_name} for the film, generate a description accordingly and make sure to include the given {field_name} in the description.\n\n"
            print(instruction)

            public_train_prompt = instruction + "\n".join(
                format_full_datum_for_prompt(field_name, datum)
                for datum in public_train_subset
            )

            if print_prompts:
                print(public_train_prompt)
                print("=========")

            if normalize_max:
                print("Exponential Mechanism")
                assert num_private_train == 0 or sigma > 0
                if num_private_train > 0:
                    # scale == sigma_calib == 1/lambda. lambda for exponential distribution.
                    sigma_calib = (2 / num_private_train_splits) * (1 / sigma)
            else:
                print("Gaussian Mechanism")
                if num_private_train_splits > 0:
                    sigma_calib = math.sqrt(2) / num_private_train_splits * sigma
                else:
                    sigma_calib = 0
            print(
                f"sigma in command {sigma}. sigma added according to sensitivity {sigma_calib}"
            )

            stop_tokens = {"\n", "<|endoftext|>", " END"}
            stop_tokens_ids = {lm.encoding.encode_single_token(t) for t in stop_tokens}

            client_session.set(aiohttp.ClientSession())

            async with client_session.get():
                for i, test_datum in enumerate(query_subset, 1):
                    print(f"# Example {i}")
                    print(f'{field_name}: "{test_datum["label"]}"')

                    np.random.seed(synth_seed + i)
                    gen_seed = np.random.randint(100000)
                    print(f"gen-seed: {gen_seed}")

                    if num_private_train_splits > 0:
                        generated_token_ids = await generate_with_private_prompts(
                            trainset,
                            num_private_train,
                            num_private_train_splits,
                            instruction,
                            public_train_prompt,
                            stop_tokens_ids,
                            test_datum["label"],
                            lm,
                            noise_rng,
                            sigma_calib,
                            field_name,
                            top_p,
                            no_public_token,
                            subsample_per_token,
                            gen_seed,
                            max_tokens=max_token_per_text
                            - 1,  # need one token length for EOS.
                            normalize_max=normalize_max,
                        )
                    else:
                        generated_token_ids = await generate_with_public_prompt(
                            public_train_prompt,
                            stop_tokens,
                            test_datum["label"],
                            lm,
                            field_name,
                            max_tokens=max_token_per_text,
                        )

                    generated = lm.encoding.decode(generated_token_ids).rstrip('"')

                    print(f"Generated: {generated}\n")
                    output_datum = {}
                    output_datum["content"] = generated.strip()
                    output_datum["label"] = test_datum["label"]
                    synthetic_examples.append(output_datum)

        if num_test > 0 and num_test <= len(data["test"]):
            test_subset = (
                data["test"]
                .shuffle(seed=12345, keep_in_memory=True)
                .select(range(num_test), keep_in_memory=True)
            )
        else:
            test_subset = data["test"]

        all_raw_answers_wout_DP = get_model_response(
            query_subset, test_subset, openai_model, field_name
        )
        all_orig_ans = []
        for resp in all_raw_answers_wout_DP:
            all_orig_ans.append(resp["text"])
        all_orig_ans = [ans.strip() for ans in all_orig_ans]
        test_labels = test_subset["label"]
        orig_accuracy = em_accuracy_helper(all_orig_ans, test_labels)
        print(f"Accuracy (original) without DP: {orig_accuracy}")

        if use_dp_prompts:
            all_raw_answers_w_DP = get_model_response(
                synthetic_examples, test_subset, openai_model, field_name
            )
            all_orig_ans = []
            for resp in all_raw_answers_w_DP:
                all_orig_ans.append(resp["text"])
            all_orig_ans = [ans.strip() for ans in all_orig_ans]
            test_labels = test_subset["label"]
            orig_accuracy = em_accuracy_helper(all_orig_ans, test_labels)
            print(f"Accuracy (original) with DP: {orig_accuracy}")

    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        traceback.print_exc()
        raise


if __name__ == "__main__":
    typer.run(_main)

src/dp_few_shot_generation/run_exp_movie.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import asyncio
import math
import re
import sys
import time
import traceback
from collections.abc import Iterable, Set
from typing import Annotated, cast

import aiohttp
import more_itertools
import numpy as np
import openai
import scipy.special
import tqdm
import typer
from datasets import DatasetDict, load_dataset
from lmapi.lm import LM, CompletionsSettings
from lmapi.openai import client_session

from dp_few_shot_generation.lm import (
    api_openai_com,
    next_logprobs,
    normalized_logprobs_for_chosen_tokens,
)
from dp_few_shot_generation.prob_utils import densify, log_max_normalize, log_normalize

DEFAULT_NUM_PRIVATE_TRAIN = 80
DEFAULT_NUM_PUBLIC_TRAIN = 0
DEFAULT_NUM_VALID = 4
DEFAULT_NUM_PRIVATE_TRAIN_SPLITS = 80
DEFAULT_NUM_TEST = 1000

labels = ["Ab", "Entity", "Description", "Person", "Location", "Number"]
label_dict = {
    0: ["Ab"],
    1: ["Entity"],
    2: ["Description"],
    3: ["Person"],
    4: ["Location"],
    5: ["Number"],
}


def format_full_datum_for_prompt(labels, datum: dict[str, str]):
    return f'Answer Type: "{labels[datum["coarse_label"]]}"\nText: "{datum["text"] + " END"}"\n'


def format_test_input_for_prompt(labels, test_input: int):
    return f'Answer Type: "{labels[test_input]}"\nText: "'


def construct_prompt_same(train_examples, test_example):
    prompt = "Classify the questions based on whether their answer type is a Number, Location, Person, Description, Entity, or Abbreviation.\n\n"
    for train_example in train_examples:
        prompt += "Question: " + train_example["text"] + "\n"
        prompt += (
            "Answer Type: " + label_dict[train_example["coarse_label"]][0] + "\n\n"
        )
    prompt += "Question: " + test_example["text"] + "\n"
    prompt += "Answer Type:"
    return prompt


def complete(prompt, l, model_name, temp=0, num_log_probs=None, echo=False, n=None):
    # call GPT-3 API until result is provided and then return it
    response = None
    received = False
    while not received:
        try:
            response = openai.Completion.create(
                engine=model_name,
                prompt=prompt,
                max_tokens=l,
                temperature=temp,
                logprobs=num_log_probs,
                echo=echo,
                stop="\n",
                n=n,
            )
            received = True
        except:
            error = sys.exc_info()[0]
            if (
                error == openai.error.InvalidRequestError
            ):  # something is wrong: e.g. prompt too long
                print(f"InvalidRequestError\nPrompt passed in:\n\n{prompt}\n\n")
                assert False

            print("API error:", error)
            time.sleep(1)
    return response


def chunks(lst, n):
    """Yield successive n-sized chunks from lst."""
    for i in range(0, len(lst), n):
        yield lst[i : i + n]


def get_model_response(data, test_examples, openai_model):
    all_raw_answers = []

    prompts = []
    train_examples = data

    for test_example in test_examples:
        prompts.append(construct_prompt_same(train_examples, test_example))

    chunked_prompts = list(chunks(prompts, 20))
    for test_chunk in chunked_prompts:
        response = complete(test_chunk, l=1, model_name=openai_model, num_log_probs=100)

        for answer_id, answer in enumerate(response["choices"]):
            all_raw_answers.append(answer)

    return all_raw_answers


def get_label_probs(all_raw_answers, test_subset):
    """Obtain model's label probability for each of the test examples. The returned prob is NOT normalized"""
    num_classes = len(label_dict)
    approx = False
    assert len(all_raw_answers) == len(test_subset)

    # Fill in the labels that is in the top k prob
    all_label_probs = []
    all_missing_positions = []
    cnt = 0
    for i, ans in enumerate(all_raw_answers):
        try:
            top_logprobs = ans["logprobs"]["top_logprobs"][
                0
            ]  # [0] since we only ask for complete one more token
        except:
            cnt += 1  # cnt for corner case
        label_probs = [0] * len(label_dict.keys())
        for j, label_list in label_dict.items():
            all_found = True
            for label in label_list:  # each possible label correspond to the same class
                label = " " + label  # notice prompt does not have space after 'A:'
                if label in top_logprobs:
                    label_probs[j] += np.exp(top_logprobs[label])
                else:
                    all_found = False
            if not all_found:
                position = (i, j)  # (which test example, which label)
                all_missing_positions.append(position)
        all_label_probs.append(label_probs)
    all_label_probs = np.array(all_label_probs)  # prob not normalized

    return all_label_probs  # NOT NORMALIZED


def eval_accuracy(all_label_probs, test_labels, mode=None, p_cf=None):
    # evaluate the accuracy with and without contextual calibration
    num_classes = all_label_probs.shape[1]
    if p_cf is None:
        # do not calibrate
        W = np.identity(num_classes)
        b = np.zeros([num_classes, 1])
    else:
        # calibrate
        if mode == "diagonal_W":
            W = np.linalg.inv(np.identity(num_classes) * p_cf)
            b = np.zeros([num_classes, 1])
        elif mode == "identity_W":
            W = np.identity(num_classes)
            b = -1 * np.expand_dims(p_cf, axis=-1)
        else:
            assert False

    correctness_list = []
    assert len(all_label_probs) == len(test_labels)
    for label_probs, true_label in zip(all_label_probs, test_labels):
        if np.sum(label_probs) > 0:  # corner case np.sum(label_probs)=0.
            label_probs = label_probs / np.sum(label_probs)  # normalize to 1

        calibrate_label_probs = np.matmul(W, np.expand_dims(label_probs, axis=-1)) + b

        ans_label = np.argmax(calibrate_label_probs)
        if ans_label == true_label:
            correctness_list.append(1)
        else:
            correctness_list.append(0)
    return np.mean(correctness_list)


def get_p_content_free(train_subset, openai_model, content_free_inputs=("N/A",)):
    """Query model with content free input, return its prediction probability for each label"""
    all_p_y = []
    for content_free_input in content_free_inputs:
        prompt = construct_prompt_same(train_subset, content_free_input)
        p_y = [0] * len(label_dict)
        for i, answers in label_dict.items():
            prob = 0
            for a in answers:
                prob += np.exp(
                    complete(
                        prompt + " " + a, 0, openai_model, echo=True, num_log_probs=1
                    )["choices"][0]["logprobs"]["token_logprobs"][-1]
                )
            p_y[i] = prob
        all_p_y.append(p_y)
    p_y = np.mean(np.array(all_p_y), axis=0)
    p_y = p_y / np.sum(p_y)  # normalize
    return p_y


def merge_logprobs_topk_mean(
    private_next_logprobs: list[dict[int, float]],
    public_next_logprobs: dict[int, float],
    n_vocab: int,
    no_public_token: bool,
    normalize_max: bool,
) -> np.ndarray:
    # Compute merged distribution
    # logsumexp - np.log(...): compute mean probability of distribution
    if normalize_max:
        normalize_func = (
            log_max_normalize  # normalize max probability to 1, Exponential mechanism
        )
    else:
        normalize_func = (
            log_normalize  # normalize sum probability to 1, Gaussian mechanism
        )
    if no_public_token:
        merged_next_logprobs = scipy.special.logsumexp(
            np.stack(
                [
                    # Turn into a 1D tensor of size n_vocab
                    densify(
                        n_vocab,
                        # Normalize distribution
                        normalize_func(
                            # Filter to the top 100 most likely next tokens according to the public prompt
                            {k: v for k, v in lps.items()}
                        ),
                    )
                    for lps in private_next_logprobs
                ]
            ),
            axis=0,
        ) - np.log(len(private_next_logprobs))

    else:
        merged_next_logprobs = scipy.special.logsumexp(
            np.stack(
                [
                    # Turn into a 1D tensor of size n_vocab
                    densify(
                        n_vocab,
                        # Normalize distribution
                        normalize_func(
                            # Filter to the top 100 most likely next tokens according to the public prompt
                            {k: v for k, v in lps.items() if k in public_next_logprobs}
                        ),
                    )
                    for lps in private_next_logprobs
                ]
            ),
            axis=0,
        ) - np.log(len(private_next_logprobs))
    merged_next_probs = np.exp(merged_next_logprobs)
    return merged_next_probs


async def generate_with_private_prompts(
    trainset,
    num_private_train,
    num_private_train_splits,
    instruction,
    public_train_prompt: str,
    stop_tokens: Set[int],
    test_input: int,
    lm: LM,
    noise_rng: np.random.RandomState,
    sigma: float,
    labels,
    top_p,
    no_public_token: bool,
    subsample_per_token: bool,
    sample_same_label_prompts: bool,
    gen_seed: int,
    max_tokens: int = 100 - 1,
    normalize_max: bool = False,
) -> list[int]:
    generated_token_ids: list[int] = []

    stringified_test_datum = format_test_input_for_prompt(labels, test_input)
    public_prompt = public_train_prompt + stringified_test_datum
    public_prompt_tokens = lm.encoding.encode(public_prompt)

    assert num_private_train_splits > 0
    if sample_same_label_prompts:
        select_list = []
        for i in range(len(trainset)):
            if trainset[i]["coarse_label"] == test_input:
                select_list.append(i)
        train_subset = trainset.select(select_list, keep_in_memory=True)
    else:
        train_subset = trainset.select(range(len(trainset)), keep_in_memory=True)

    if not subsample_per_token:
        private_train_subset = cast(
            Iterable[dict[str, str]],
            train_subset.shuffle(gen_seed, keep_in_memory=True).select(
                range(num_private_train), keep_in_memory=True
            ),
        )
        private_train_splits = [
            list(it)
            for it in more_itertools.distribute(
                num_private_train_splits, private_train_subset
            )
        ]
        private_train_prompts = [
            instruction
            + "\n".join(format_full_datum_for_prompt(labels, datum) for datum in split)
            for split in private_train_splits
        ]
        private_prompts = [
            train_prompt + "\n" + stringified_test_datum
            for train_prompt in private_train_prompts
        ]
        private_prompts_tokens = [
            lm.encoding.encode(prompt) for prompt in private_prompts
        ]

    cnt = 0
    for _ in tqdm.tqdm(range(max_tokens), total=float("inf"), unit=" tokens generated"):
        private_next_logprobs: list[dict[int, float]]
        public_next_logprobs: dict[int, float]
        # Split training dataset
        if subsample_per_token:
            private_train_subset = cast(
                Iterable[dict[str, str]],
                train_subset.shuffle(gen_seed + cnt, keep_in_memory=True).select(
                    range(num_private_train), keep_in_memory=True
                ),
            )
            cnt += 1
            private_train_splits = [
                list(it)
                for it in more_itertools.distribute(
                    num_private_train_splits, private_train_subset
                )
            ]
            # Turn the data into prompts
            private_train_prompts = [
                instruction
                + "\n".join(
                    format_full_datum_for_prompt(labels, datum) for datum in split
                )
                for split in private_train_splits
            ]
            private_prompts = [
                train_prompt + "\n" + stringified_test_datum
                for train_prompt in private_train_prompts
            ]
            private_prompts_tokens = [
                lm.encoding.encode(prompt) for prompt in private_prompts
            ]
        if no_public_token:
            private_next_logprobs = await asyncio.gather(
                *(
                    next_logprobs(lm, prompt + generated_token_ids, top_p=top_p)
                    for prompt in private_prompts_tokens
                )
            )
            merged_next_probs = merge_logprobs_topk_mean(
                private_next_logprobs,
                None,
                lm.encoding.n_vocab,
                no_public_token,
                normalize_max,
            )
            if normalize_max:
                # scale = 1/lambda
                noise = noise_rng.exponential(scale=sigma, size=lm.encoding.n_vocab)
            else:
                noise = noise_rng.normal(0, sigma, size=lm.encoding.n_vocab)
            merged_next_probs += noise
        else:
            public_next_logprobs = await next_logprobs(
                lm, public_prompt_tokens + generated_token_ids, top_p=top_p
            )
            private_next_logprobs = await asyncio.gather(
                *(
                    normalized_logprobs_for_chosen_tokens(
                        lm,
                        prompt + generated_token_ids,
                        public_next_logprobs.keys(),
                        top_p=top_p,
                    )
                    for prompt in private_prompts_tokens
                )
            )
            merged_next_probs = merge_logprobs_topk_mean(
                private_next_logprobs,
                public_next_logprobs,
                lm.encoding.n_vocab,
                no_public_token,
                normalize_max,
            )
            if normalize_max:
                # scale = 1/lambda
                noise = noise_rng.exponential(
                    scale=sigma, size=len(public_next_logprobs)
                )
            else:
                noise = noise_rng.normal(0, sigma, size=len(public_next_logprobs))
            merged_next_probs[list(public_next_logprobs.keys())] += noise

        next_token_id = int(np.argmax(merged_next_probs))

        if next_token_id in stop_tokens:
            break

        generated_token_ids.append(next_token_id)

        del next_token_id
    return generated_token_ids


async def generate_with_public_prompt(
    public_train_prompt: str,
    stop_tokens: Set[str],
    test_input: str,
    lm: LM,
    labels,
    max_tokens: int = 500,
) -> list[int]:
    public_prompt = public_train_prompt + format_test_input_for_prompt(
        labels, test_input
    )
    public_prompt_tokens = lm.encoding.encode(public_prompt)
    public_prompt_tokens = public_prompt

    [completion] = await lm.completions(
        public_prompt_tokens,
        CompletionsSettings(
            temperature=0.0, max_tokens=max_tokens, n=1, stop=list(stop_tokens)
        ),
    )
    generated_tokens = [st.token.token_id for st in completion]
    return generated_tokens


def select_uniform_n_shots_over_labels(data, n_shots):
    select_list = []
    n_shots_per_label = math.ceil(n_shots / len(labels))
    labels_counter = {label[1][0]: n_shots_per_label for label in label_dict.items()}
    n_shots_selected = 0
    for i in range(len(data)):
        label = label_dict[data[i]["coarse_label"]][0]
        if labels_counter[label] == 0 or data[i]["coarse_label"] == 0:
            continue
        else:
            labels_counter[label] -= 1
            select_list.append(i)
            n_shots_selected += 1
        if n_shots_selected == n_shots:
            break
    query_subset = data.select(select_list, keep_in_memory=True)
    return query_subset


def _main(
    sigma: Annotated[float, typer.Option()],  # noise parameters
    openai_model: Annotated[str, typer.Option()] = "babbage",
    print_prompts: Annotated[bool, typer.Option()] = False,
    # num_private_train=MN. MN=0 with num_valid=4 will get epsilon=0 (4-shot) results.
    num_private_train: Annotated[int, typer.Option()] = DEFAULT_NUM_PRIVATE_TRAIN,
    # by default set to 0. set_num_public_train >0 indicates additional public data available.
    set_num_public_train: Annotated[int, typer.Option()] = DEFAULT_NUM_PUBLIC_TRAIN,
    # num_valid=n. n samples to be generated for n-shot ICL
    num_valid: Annotated[int, typer.Option()] = DEFAULT_NUM_VALID,
    # num_private_train_splits=M
    num_private_train_splits: Annotated[
        int, typer.Option()
    ] = DEFAULT_NUM_PRIVATE_TRAIN_SPLITS,
    num_test: Annotated[int, typer.Option()] = DEFAULT_NUM_TEST,
    # no_public_token=True, RVP=False; no_public_token=False, RVP=True
    no_public_token: Annotated[bool, typer.Option()] = False,
    # subsample_per_token=True: at each token generation, subsample a new test set
    subsample_per_token: Annotated[bool, typer.Option()] = False,
    use_dp_prompts: Annotated[bool, typer.Option()] = False,
    # sample_same_label_prompts=True: sample subsets from the sets with same targeted labels.
    sample_same_label_prompts: Annotated[bool, typer.Option()] = False,
    # normalize_max=True, Exponential mechanism; normalize_max=False, Gaussian mechanism
    normalize_max: Annotated[bool, typer.Option()] = False,
    # max_token_per_text=T_max
    max_token_per_text: Annotated[int, typer.Option()] = 15,
    # consistent with default parameters in the documentation https://learn.microsoft.com/en-us/azure/cognitive-services/openai/reference#completions
    top_p: Annotated[float, typer.Option()] = 1,
    # random seed for subsampling in generation
    synth_seed: Annotated[int, typer.Option()] = 0,
    # random seed for n-shot demonstrations sampling in evaluation
    eval_seed: Annotated[int, typer.Option()] = 0,
):
    async def main():
        if (num_private_train == 0) != (num_private_train_splits == 0):
            raise ValueError(
                "Either both or neither of --num-private-train and --num-private-train-splits can be 0"
            )
        command = ["python", sys.argv[0]]
        for x in sys.argv[1:]:
            if x.startswith("--"):
                assert '"' not in x and "'" not in x
                command.append(x)
            else:
                assert "'" not in x
                if re.match("^[a-zA-Z0-9_]+$", x):
                    command.append("%s" % x)
                else:
                    command.append("'%s'" % x)
        command = " ".join(command)
        print(command)

        if no_public_token:
            num_public_train = 0
        else:
            num_public_train = set_num_public_train

        lm = api_openai_com(openai_model)
        noise_rng = np.random.RandomState()

        data = cast(DatasetDict, load_dataset("trec"))
        print(labels)

        trainset = data["train"].shuffle(seed=synth_seed, keep_in_memory=True)
        print("trainset length", len(trainset))
        if num_public_train > 0:
            public_train_subset = cast(
                Iterable[dict[str, str]],
                trainset.select(
                    range(
                        len(trainset) - num_public_train,
                        len(trainset),
                        keep_in_memory=True,
                    )
                ),
            )
        else:
            public_train_subset = []

        trainset = trainset.select(
            range(len(trainset) - num_public_train), keep_in_memory=True
        )
        queryset = data["train"].shuffle(seed=eval_seed, keep_in_memory=True)
        query_subset = select_uniform_n_shots_over_labels(queryset, num_valid)

        if use_dp_prompts:
            synthetic_examples = []

            # Turn the data into prompts
            instruction = "Given a label of answer type, generate a question based on the given answer type accordingly.\n\n"

            public_train_prompt = instruction + "\n".join(
                format_full_datum_for_prompt(labels, datum)
                for datum in public_train_subset
            )

            if print_prompts:
                print(public_train_prompt)
                print("=========")

            if normalize_max:
                print("Exponential Mechanism")
                assert num_private_train == 0 or sigma > 0
                if num_private_train > 0:
                    # scale == sigma_calib == 1/lambda. lambda for exponential distribution.
                    sigma_calib = (2 / num_private_train_splits) * (1 / sigma)
            else:
                print("Gaussian Mechanism")
                if num_private_train_splits > 0:
                    sigma_calib = math.sqrt(2) / num_private_train_splits * sigma
                else:
                    sigma_calib = 0
            print(
                f"sigma in command {sigma}. sigma added according to sensitivity {sigma_calib}"
            )

            stop_tokens = {"\n", "<|endoftext|>", " END"}
            stop_tokens_ids = {lm.encoding.encode_single_token(t) for t in stop_tokens}

            client_session.set(aiohttp.ClientSession())

            len_token = []
            async with client_session.get():
                for i, test_datum in enumerate(query_subset, 1):
                    print(f"# Example {i}")
                    print(f'Answer Type: "{labels[test_datum["coarse_label"]]}"')
                    print(f'References:\n "{test_datum["text"]}"')

                    np.random.seed(synth_seed + i)
                    gen_seed = np.random.randint(100000)
                    print(f"gen-seed: {gen_seed}")

                    if num_private_train_splits > 0:
                        generated_token_ids = await generate_with_private_prompts(
                            trainset,
                            num_private_train,
                            num_private_train_splits,
                            instruction,
                            public_train_prompt,
                            stop_tokens_ids,
                            test_datum["coarse_label"],
                            lm,
                            noise_rng,
                            sigma_calib,
                            labels,
                            top_p,
                            no_public_token,
                            subsample_per_token,
                            sample_same_label_prompts,
                            gen_seed,
                            max_tokens=max_token_per_text
                            - 1,  # need one token length for EOS.
                            normalize_max=normalize_max,
                        )
                    else:
                        generated_token_ids = await generate_with_public_prompt(
                            public_train_prompt,
                            stop_tokens,
                            test_datum["coarse_label"],
                            lm,
                            labels,
                            max_tokens=max_token_per_text,
                        )

                    generated = lm.encoding.decode(generated_token_ids).rstrip('"')

                    print(f"Generated: {generated}\n")
                    output_datum = {}
                    output_datum["text"] = generated.strip()
                    output_datum["coarse_label"] = test_datum["coarse_label"]
                    synthetic_examples.append(output_datum)

        if num_test > 0:
            test_subset = data["test"]
            test_labels = [test_example["coarse_label"] for test_example in test_subset]

            content_free_inputs = [{"text": "N/A"}, {"text": ""}, {"text": "[MASK]"}]
            p_cf_wout_DP = get_p_content_free(
                query_subset, openai_model, content_free_inputs=content_free_inputs
            )

            all_raw_answers_wout_DP = get_model_response(
                query_subset, test_subset, openai_model
            )
            all_label_probs_wout_DP = get_label_probs(
                all_raw_answers_wout_DP, test_subset
            )

            acc_original_wout_DP = eval_accuracy(all_label_probs_wout_DP, test_labels)
            acc_calibrated_wout_DP = eval_accuracy(
                all_label_probs_wout_DP,
                test_labels,
                mode="diagonal_W",
                p_cf=p_cf_wout_DP,
            )

            print(f"Accuracy (original) without DP: {acc_original_wout_DP}")
            print(f"Accuracy (calibrated) without DP: {acc_calibrated_wout_DP}")

            if use_dp_prompts:
                p_cf_w_DP = get_p_content_free(
                    synthetic_examples,
                    openai_model,
                    content_free_inputs=content_free_inputs,
                )

                all_raw_answers_w_DP = get_model_response(
                    synthetic_examples, test_subset, openai_model
                )
                all_label_probs_w_DP = get_label_probs(
                    all_raw_answers_w_DP, test_subset
                )

                acc_original_w_DP = eval_accuracy(all_label_probs_w_DP, test_labels)
                acc_calibrated_w_DP = eval_accuracy(
                    all_label_probs_w_DP, test_labels, mode="diagonal_W", p_cf=p_cf_w_DP
                )

                print(f"Accuracy (original) with DP: {acc_original_w_DP}")
                print(f"Accuracy (calibrated) with DP: {acc_calibrated_w_DP}")

    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        traceback.print_exc()
        raise


if __name__ == "__main__":
    typer.run(_main)

src/dp_few_shot_generation/run_exp_trec.py
microsoft/dp-few-shot-generation
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

# ruff: noqa: F401
def test_package_import():
    "Tests that dp_few_shot_generation is importable."
    import dp_few_shot_generation

tests/test_dp_few_shot_generation/test_import.py
microsoft/dp-few-shot-generation