diff --git a/py/autoevals/llm.py b/py/autoevals/llm.py index c85a8a3..26bebde 100644 --- a/py/autoevals/llm.py +++ b/py/autoevals/llm.py @@ -79,6 +79,7 @@ def __init__( temperature=None, engine=None, api_key=None, + base_url=None, ): self.name = name self.model = model @@ -92,6 +93,8 @@ def __init__( self.extra_args["max_tokens"] = max(max_tokens, 5) if api_key: self.extra_args["api_key"] = api_key + if base_url: + self.extra_args["base_url"] = base_url self.render_args = {} if render_args: @@ -199,6 +202,7 @@ def __init__( temperature=0, engine=None, api_key=None, + base_url=None, ): choice_strings = list(choice_scores.keys()) @@ -220,6 +224,7 @@ def __init__( temperature=temperature, engine=engine, api_key=api_key, + base_url=base_url, render_args={"__choices": choice_strings}, ) @@ -235,7 +240,7 @@ def from_spec_file(cls, name: str, path: str, **kwargs): class SpecFileClassifier(LLMClassifier): - def __new__(cls, model=None, engine=None, use_cot=None, max_tokens=None, temperature=None, api_key=None): + def __new__(cls, model=None, engine=None, use_cot=None, max_tokens=None, temperature=None, api_key=None, base_url=None): kwargs = {} if model is not None: kwargs["model"] = model @@ -249,6 +254,8 @@ def __new__(cls, model=None, engine=None, use_cot=None, max_tokens=None, tempera kwargs["temperature"] = temperature if api_key is not None: kwargs["api_key"] = api_key + if base_url is not None: + kwargs["base_url"] = base_url # convert FooBar to foo_bar template_name = re.sub(r"(?