Merge pull request stanfordnlp#1063 from Neoxelox/main

Allow to configure Anyscale API in constructor
inayet · May 28, 2024 · 78d95df · 78d95df
2 parents 41d0f0e + 57ea58c
commit 78d95df
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/dsp/modules/hf_client.py b/dsp/modules/hf_client.py
@@ -397,11 +397,12 @@ def _generate(self, prompt, use_chat_api=False, **kwargs):
 
 
 class Anyscale(HFModel):
-    def __init__(self, model, **kwargs):
+    def __init__(self, model, api_base="https://api.endpoints.anyscale.com/v1", api_key=None, **kwargs):
         super().__init__(model=model, is_client=True)
         self.session = requests.Session()
-        self.api_base = os.getenv("ANYSCALE_API_BASE")
-        self.token = os.getenv("ANYSCALE_API_KEY")
+        self.api_base = os.getenv("ANYSCALE_API_BASE") or api_base
+        assert not self.api_base.endswith("/"), "Anyscale base URL shouldn't end with /"
+        self.token = os.getenv("ANYSCALE_API_KEY") or api_key
         self.model = model
         self.kwargs = {
             "temperature": 0.0,