diff --git a/README.md b/README.md index 0a0ac114..12b31125 100644 --- a/README.md +++ b/README.md @@ -86,6 +86,14 @@ For the [Anthropic API](https://docs.anthropic.com/claude/docs/quickstart-guide) pip install anthropic ``` +For the [xAI API](https://docs.x.ai/docs/overview): +```bash +pip install openai + +# Create a .env file and add your XAI API key +echo "XAI_API_KEY=" > .env +``` + ## Running the evals ```bash python -m simple-evals.simple_evals --list-models diff --git a/sampler/grok_sampler.py b/sampler/grok_sampler.py new file mode 100644 index 00000000..603f89d1 --- /dev/null +++ b/sampler/grok_sampler.py @@ -0,0 +1,59 @@ +import os +import time +from typing import Any +from openai import OpenAI + +from ..types import MessageList, SamplerBase + +GROK_SYSTEM_MESSAGE = "You are Grok 2, an AI built by xAI, designed to answer almost any question with an outside perspective on humanity, aiming for maximum helpfulness." + + +class GrokCompletionSampler(SamplerBase): + """ + Sample from X.AI's Grok model via their API + """ + + def __init__( + self, + model: str = "grok-2-1212", + system_message: str | None = None, + temperature: float = 0.5, + max_tokens: int = 1024, + ): + self.api_key_name = "XAI_API_KEY" + self.client = OpenAI( + api_key=os.environ.get("XAI_API_KEY"), + base_url="https://api.x.ai/v1", + ) + self.model = model + self.system_message = system_message + self.temperature = temperature + self.max_tokens = max_tokens + + def _pack_message(self, role: str, content: Any): + return {"role": str(role), "content": content} + + def __call__(self, message_list: MessageList) -> str: + if self.system_message: + message_list = [ + self._pack_message("system", self.system_message) + ] + message_list + trial = 0 + while True: + try: + response = self.client.chat.completions.create( + model=self.model, + messages=message_list, + temperature=self.temperature, + max_tokens=self.max_tokens, + ) + return response.choices[0].message.content + except Exception as e: + exception_backoff = 2**trial # exponential back off + print( + f"Rate limit exception so wait and retry {trial} after {exception_backoff} sec", + e, + ) + time.sleep(exception_backoff) + trial += 1 + # unknown error shall throw exception diff --git a/simple_evals.py b/simple_evals.py index e3debd3b..476a1d77 100644 --- a/simple_evals.py +++ b/simple_evals.py @@ -16,6 +16,7 @@ ) from .sampler.o1_chat_completion_sampler import O1ChatCompletionSampler from .sampler.claude_sampler import ClaudeCompletionSampler, CLAUDE_SYSTEM_MESSAGE_LMSYS +from .sampler.grok_sampler import GrokCompletionSampler, GROK_SYSTEM_MESSAGE def main(): @@ -47,7 +48,7 @@ def main(): ), "o1": O1ChatCompletionSampler( model="o1", - ), + ), "o1-preview": O1ChatCompletionSampler( model="o1-preview", ), @@ -82,6 +83,11 @@ def main(): model="claude-3-opus-20240229", system_message=CLAUDE_SYSTEM_MESSAGE_LMSYS, ), + # grok models: + "grok-2-1212": GrokCompletionSampler( + model="grok-2-1212", + system_message=GROK_SYSTEM_MESSAGE, + ), } if args.list_models: