Grok Sampler

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=<your-api-key>" > .env
+```
+
 ## Running the evals
 ```bash
 python -m simple-evals.simple_evals --list-models

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

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: