Skip to content

Commit

Permalink
Fix ClaudeModel support.
Browse files Browse the repository at this point in the history
  • Loading branch information
shaoyijia committed Apr 26, 2024
1 parent 7487b90 commit 853c188
Showing 1 changed file with 12 additions and 4 deletions.
16 changes: 12 additions & 4 deletions src/lm.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,11 @@
from dsp.modules.hf_client import send_hfvllm_request_v00, send_hftgi_request_v01_wrapped
from transformers import AutoTokenizer

try:
import anthropic
except ImportError:
pass


class OpenAIModel(dspy.OpenAI):
"""A wrapper class for dspy.OpenAI."""
Expand Down Expand Up @@ -115,7 +120,7 @@ def __init__(
):
super().__init__(model)
try:
from anthropic import Anthropic, RateLimitError
from anthropic import Anthropic
except ImportError as err:
raise ImportError("Claude requires `pip install anthropic`.") from err

Expand Down Expand Up @@ -188,7 +193,7 @@ def basic_request(self, prompt: str, **kwargs):

@backoff.on_exception(
backoff.expo,
(RateLimitError,),
(anthropic.RateLimitError,),
max_time=1000,
max_tries=8,
on_backoff=backoff_hdlr,
Expand Down Expand Up @@ -219,8 +224,11 @@ def __call__(self, prompt, only_completed=True, return_sorted=False, **kwargs):
for _ in range(n):
response = self.request(prompt, **kwargs)
self.log_usage(response)
if only_completed and response.stop_reason == "max_tokens":
continue
# This is the original behavior in dspy/dsp/modules/anthropic.py.
# Comment it out because it can cause "IndexError: list index out of range" silently
# which is not transparent to developers.
# if only_completed and response.stop_reason == "max_tokens":
# continue
completions = [c.text for c in response.content]
return completions

Expand Down

0 comments on commit 853c188

Please sign in to comment.