Skip to content

Commit

Permalink
Add chat client for Palmyra (#2958)
Browse files Browse the repository at this point in the history
  • Loading branch information
yifanmai authored Aug 29, 2024
1 parent bbdb00e commit be2fbd4
Showing 1 changed file with 25 additions and 0 deletions.
25 changes: 25 additions & 0 deletions src/helm/clients/palmyra_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import requests
from typing import Any, Dict, List

from helm.clients.openai_client import OpenAIClient
from helm.common.cache import CacheConfig
from helm.common.hierarchical_logger import hlog
from helm.common.request import wrap_request_time, Request, RequestResult, GeneratedOutput, Token, ErrorFlags
Expand Down Expand Up @@ -142,3 +143,27 @@ def do_it() -> Dict[str, Any]:
completions=completions,
embedding=[],
)


class PalmyraChatClient(OpenAIClient):
"""Sends request to a Palmyra model using a OpenAI-compatible Chat API."""

def __init__(
self,
tokenizer: Tokenizer,
tokenizer_name: str,
cache_config: CacheConfig,
api_key: str,
base_url: str,
):
super().__init__(
tokenizer=tokenizer,
tokenizer_name=tokenizer_name,
cache_config=cache_config,
api_key=api_key,
org_id=None,
base_url=base_url,
)

def _is_chat_model_engine(self, model_engine: str) -> bool:
return True

0 comments on commit be2fbd4

Please sign in to comment.