code-refactor

explosion · Mar 29, 2024 · 6ed4ca4 · 6ed4ca4
1 parent 0555759
commit 6ed4ca4
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/spacy_llm/models/rest/mistral/model.py b/spacy_llm/models/rest/mistral/model.py
@@ -3,6 +3,9 @@
 from typing import Iterable, Optional, Any, Dict
 from ..base import REST
 
+from mistralai.client import MistralClient
+from mistralai.models.chat_completion import ChatMessage
+
 
 class AzureMistral(REST):
     def __init__(
@@ -29,13 +32,10 @@ def __init__(
 
     def __call__(self, prompts: Iterable[Iterable[str]]) -> Iterable[Iterable[str]]:
         all_resps = []
+        api_key = self._credentials.get("api-key")
         for prompts_doc in prompts:
             doc_resps = []
             for prompt in prompts_doc:
-                from mistralai.client import MistralClient
-                from mistralai.models.chat_completion import ChatMessage
-
-                api_key = self._credentials.get("api-key")
                 client = MistralClient(endpoint=self._endpoint, api_key=api_key)
 
                 chat_response = client.chat(