anthropic_llms.py 5.79 KB
Newer Older
Jason Phang's avatar
Jason Phang committed
1
import os
haileyschoelkopf's avatar
haileyschoelkopf committed
2
3
from lm_eval.api.model import LM
from lm_eval.api.registry import register_model
Jason Phang's avatar
Jason Phang committed
4
5
from tqdm import tqdm
import time
baberabb's avatar
baberabb committed
6
from lm_eval.logger import eval_logger
baberabb's avatar
baberabb committed
7
from typing import List, Literal, Any
Jason Phang's avatar
Jason Phang committed
8
9


lintangsutawika's avatar
lintangsutawika committed
10
def anthropic_completion(
11
    client,  #: anthropic.Anthropic,
baberabb's avatar
baberabb committed
12
13
14
15
16
    model: str,
    prompt: str,
    max_tokens_to_sample: int,
    temperature: float,
    stop: List[str],
baberabb's avatar
baberabb committed
17
    **kwargs: Any,
lintangsutawika's avatar
lintangsutawika committed
18
):
Jason Phang's avatar
Jason Phang committed
19
20
21
22
    """Query Anthropic API for completion.

    Retry with back-off until they respond
    """
23
24
25
26
27
28
29
30
31

    try:
        import anthropic
    except ModuleNotFoundError:
        raise Exception(
            "attempted to use 'anthropic' LM type, but package `anthropic` is not installed. \
please install anthropic via `pip install lm-eval[anthropic]` or `pip install -e .[anthropic]`",
        )

Jason Phang's avatar
Jason Phang committed
32
33
34
    backoff_time = 3
    while True:
        try:
baberabb's avatar
baberabb committed
35
            response = client.completions.create(
Jason Phang's avatar
Jason Phang committed
36
37
38
39
40
41
42
                prompt=f"{anthropic.HUMAN_PROMPT} {prompt}{anthropic.AI_PROMPT}",
                model=model,
                # NOTE: Claude really likes to do CoT, and overly aggressive stop sequences
                #       (e.g. gsm8k's ":") may truncate a lot of the input.
                stop_sequences=[anthropic.HUMAN_PROMPT] + stop,
                max_tokens_to_sample=max_tokens_to_sample,
                temperature=temperature,
baberabb's avatar
baberabb committed
43
                **kwargs,
Jason Phang's avatar
Jason Phang committed
44
            )
baberabb's avatar
baberabb committed
45
46
47
48
49
            return response.completion
        except anthropic.RateLimitError as e:
            eval_logger.warning(
                f"RateLimitError occurred: {e.__cause__}\n Retrying in {backoff_time} seconds"
            )
Jason Phang's avatar
Jason Phang committed
50
51
52
53
            time.sleep(backoff_time)
            backoff_time *= 1.5


haileyschoelkopf's avatar
haileyschoelkopf committed
54
@register_model("anthropic")
lintangsutawika's avatar
lintangsutawika committed
55
class AnthropicLM(LM):
baberabb's avatar
baberabb committed
56
    REQ_CHUNK_SIZE = 20  # TODO: not used
Jason Phang's avatar
Jason Phang committed
57

baberabb's avatar
baberabb committed
58
59
    def __init__(
        self,
60
        batch_size: int = 1,
baberabb's avatar
baberabb committed
61
62
        model: str = "claude-2.0",
        max_tokens_to_sample: int = 256,
63
64
65
        temperature: float = 0,  # defaults to 1
        **kwargs,  # top_p, top_k, etc.
    ):
baberabb's avatar
baberabb committed
66
        """Anthropic API wrapper.
Jason Phang's avatar
Jason Phang committed
67
68

        :param model: str
baberabb's avatar
baberabb committed
69
            Anthropic model e.g. 'claude-instant-v1', 'claude-2'
baberabb's avatar
baberabb committed
70
71
72
73
74
75
        :param max_tokens_to_sample: int
            Maximum number of tokens to sample from the model
        :param temperature: float
            Sampling temperature
        :param kwargs: Any
            Additional model_args to pass to the API client
Jason Phang's avatar
Jason Phang committed
76
77
        """
        super().__init__()
lintangsutawika's avatar
lintangsutawika committed
78

79
80
81
82
83
84
85
86
        try:
            import anthropic
        except ModuleNotFoundError:
            raise Exception(
                "attempted to use 'anthropic' LM type, but package `anthropic` is not installed. \
please install anthropic via `pip install lm-eval[anthropic]` or `pip install -e .[anthropic]`",
            )

Jason Phang's avatar
Jason Phang committed
87
        self.model = model
baberabb's avatar
baberabb committed
88
        # defaults to os.environ.get("ANTHROPIC_API_KEY")
baberabb's avatar
baberabb committed
89
        self.client = anthropic.Anthropic()
baberabb's avatar
baberabb committed
90
91
92
        self.temperature = temperature
        self.max_tokens_to_sample = max_tokens_to_sample
        self.tokenizer = self.client.get_tokenizer()
baberabb's avatar
baberabb committed
93
        self.kwargs = kwargs
Jason Phang's avatar
Jason Phang committed
94
95
96

    @property
    def eot_token_id(self):
baberabb's avatar
baberabb committed
97
        # Not sure but anthropic.AI_PROMPT -> [203, 203, 50803, 30]
Jason Phang's avatar
Jason Phang committed
98
99
100
101
102
103
104
105
        raise NotImplementedError("No idea about anthropic tokenization.")

    @property
    def max_length(self):
        return 2048

    @property
    def max_gen_toks(self):
baberabb's avatar
baberabb committed
106
        return self.max_tokens_to_sample
Jason Phang's avatar
Jason Phang committed
107
108
109
110

    @property
    def batch_size(self):
        # Isn't used because we override _loglikelihood_tokens
baberabb's avatar
baberabb committed
111
        raise NotImplementedError("No support for logits.")
Jason Phang's avatar
Jason Phang committed
112
113
114
115

    @property
    def device(self):
        # Isn't used because we override _loglikelihood_tokens
baberabb's avatar
baberabb committed
116
        raise NotImplementedError("No support for logits.")
Jason Phang's avatar
Jason Phang committed
117

baberabb's avatar
baberabb committed
118
119
    def tok_encode(self, string: str) -> List[int]:
        return self.tokenizer.encode(string).ids
Jason Phang's avatar
Jason Phang committed
120

baberabb's avatar
baberabb committed
121
122
    def tok_decode(self, tokens: List[int]) -> str:
        return self.tokenizer.decode(tokens)
Jason Phang's avatar
Jason Phang committed
123
124
125
126
127
128
129
130

    def _loglikelihood_tokens(self, requests, disable_tqdm=False):
        raise NotImplementedError("No support for logits.")

    def greedy_until(self, requests):
        if not requests:
            return []

haileyschoelkopf's avatar
haileyschoelkopf committed
131
132
        requests = [req.args for req in requests]

Jason Phang's avatar
Jason Phang committed
133
134
        res = []
        for request in tqdm(requests):
baberabb's avatar
baberabb committed
135
136
137
            try:
                inp = request[0]
                request_args = request[1]
138
139
140
141
                # generation_kwargs
                until = request_args.get("until")
                max_gen_toks = request_args.get("max_gen_toks", self.max_length)
                temperature = request_args.get("temperature", self.temperature)
baberabb's avatar
baberabb committed
142
143
144
145
                response = anthropic_completion(
                    client=self.client,
                    model=self.model,
                    prompt=inp,
146
147
                    max_tokens_to_sample=max_gen_toks,
                    temperature=temperature,  # TODO: implement non-greedy sampling for Anthropic
baberabb's avatar
baberabb committed
148
149
150
151
152
153
                    stop=until,
                    **self.kwargs,
                )
                res.append(response)

                self.cache_hook.add_partial("greedy_until", request, response)
154
            except anthropic.APIConnectionError as e:  # noqa: F821
baberabb's avatar
baberabb committed
155
156
                eval_logger.critical(f"Server unreachable: {e.__cause__}")
                break
157
            except anthropic.APIStatusError as e:  # noqa: F821
baberabb's avatar
baberabb committed
158
159
                eval_logger.critical(f"API error {e.status_code}: {e.message}")
                break
haileyschoelkopf's avatar
haileyschoelkopf committed
160

Jason Phang's avatar
Jason Phang committed
161
162
163
164
165
166
167
168
169
        return res

    def _model_call(self, inps):
        # Isn't used because we override _loglikelihood_tokens
        raise NotImplementedError()

    def _model_generate(self, context, max_length, eos_token_id):
        # Isn't used because we override greedy_until
        raise NotImplementedError()
baberabb's avatar
baberabb committed
170
171
172
173
174
175

    def loglikelihood(self, requests):
        raise NotImplementedError("No support for logits.")

    def loglikelihood_rolling(self, requests):
        raise NotImplementedError("No support for logits.")