Merge pull request #431 from explosion/main

Sync `develop` with `main`
explosion · Jan 29, 2024 · 89c7971 · 89c7971
2 parents ad30add + 0fc4633
commit 89c7971
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 22 deletions.
diff --git a/spacy_llm/models/rest/azure/model.py b/spacy_llm/models/rest/azure/model.py
@@ -35,7 +35,7 @@ def __init__(
         self._deployment_name = deployment_name
         super().__init__(
             name=name,
-            endpoint=endpoint or endpoint,
+            endpoint=endpoint,
             config=config,
             strict=strict,
             max_tries=max_tries,

diff --git a/spacy_llm/models/rest/openai/model.py b/spacy_llm/models/rest/openai/model.py
@@ -36,12 +36,6 @@ def credentials(self) -> Dict[str, str]:
         if api_org:
             headers["OpenAI-Organization"] = api_org
 
-        # Ensure endpoint is supported.
-        if self._endpoint not in (Endpoints.NON_CHAT, Endpoints.CHAT):
-            raise ValueError(
-                f"Endpoint {self._endpoint} isn't supported. Please use one of: {Endpoints.CHAT}, {Endpoints.NON_CHAT}."
-            )
-
         return headers
 
     def _verify_auth(self) -> None:
@@ -115,9 +109,21 @@ def _request(json_data: Dict[str, Any]) -> Dict[str, Any]:
 
                 return responses
 
-            if self._endpoint == Endpoints.CHAT:
-                # The OpenAI API doesn't support batching for /chat/completions yet, so we have to send individual
-                # requests.
+            # The OpenAI API doesn't support batching for /chat/completions yet, so we have to send individual requests.
+
+            if self._endpoint == Endpoints.NON_CHAT:
+                responses = _request({"prompt": prompts_for_doc})
+                if "error" in responses:
+                    return responses["error"]
+                assert len(responses["choices"]) == len(prompts_for_doc)
+
+                for response in responses["choices"]:
+                    if "text" in response:
+                        api_responses.append(response["text"])
+                    else:
+                        api_responses.append(srsly.json_dumps(response))
+
+            else:
                 for prompt in prompts_for_doc:
                     responses = _request(
                         {"messages": [{"role": "user", "content": prompt}]}
@@ -134,18 +140,6 @@ def _request(json_data: Dict[str, Any]) -> Dict[str, Any]:
                         )
                     )
 
-            elif self._endpoint == Endpoints.NON_CHAT:
-                responses = _request({"prompt": prompts_for_doc})
-                if "error" in responses:
-                    return responses["error"]
-                assert len(responses["choices"]) == len(prompts_for_doc)
-
-                for response in responses["choices"]:
-                    if "text" in response:
-                        api_responses.append(response["text"])
-                    else:
-                        api_responses.append(srsly.json_dumps(response))
-
             all_api_responses.append(api_responses)
 
         return all_api_responses

diff --git a/spacy_llm/tests/tasks/legacy/test_ner.py b/spacy_llm/tests/tasks/legacy/test_ner.py
@@ -832,6 +832,7 @@ def test_ner_to_disk(noop_config, tmp_path: Path):
     assert task1._label_dict == task2._label_dict == labels
 
 
+@pytest.mark.filterwarnings("ignore:Task supports sharding")
 def test_label_inconsistency():
     """Test whether inconsistency between specified labels and labels in examples is detected."""
     cfg = f"""

diff --git a/spacy_llm/tests/tasks/test_ner.py b/spacy_llm/tests/tasks/test_ner.py
@@ -820,6 +820,7 @@ def test_ner_to_disk(noop_config: str, tmp_path: Path):
     assert task1._label_dict == task2._label_dict == labels
 
 
+@pytest.mark.filterwarnings("ignore:Task supports sharding")
 def test_label_inconsistency():
     """Test whether inconsistency between specified labels and labels in examples is detected."""
     cfg = f"""