84 lines
2.7 KiB
Python
84 lines
2.7 KiB
Python
from typing import Iterator, Optional, Callable
|
|
from mistralai import Mistral
|
|
from mistral_common.protocol.instruct.messages import SystemMessage, UserMessage
|
|
from mistral_common.protocol.instruct.request import ChatCompletionRequest
|
|
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
|
|
|
|
from . import LlmEngine
|
|
from ..util import skip_prefix
|
|
|
|
class MistralLlmEngine(LlmEngine):
|
|
def __init__(
|
|
self,
|
|
model: str,
|
|
temperature: float,
|
|
token_limit: int,
|
|
api_key: str,
|
|
):
|
|
self._model = model
|
|
self._temperature = temperature
|
|
self._token_limit = token_limit
|
|
self._api_key = api_key
|
|
self._client = Mistral(api_key=api_key)
|
|
self._tokenizer = MistralTokenizer.v3()
|
|
|
|
def infer(self, system_prompt: str, main_context: str, continuation_text: str, should_stop: Callable[[], bool] = lambda: False) -> Iterator[str]:
|
|
messages = [
|
|
{
|
|
"role": "system",
|
|
"content": system_prompt,
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": main_context,
|
|
},
|
|
{
|
|
"role": "assistant",
|
|
"content": continuation_text,
|
|
"prefix": True,
|
|
},
|
|
] if continuation_text else [
|
|
{
|
|
"role": "system",
|
|
"content": system_prompt,
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": main_context,
|
|
},
|
|
]
|
|
stream_response = self._client.chat.stream(
|
|
model=self._model,
|
|
messages=messages,
|
|
temperature=self._temperature,
|
|
)
|
|
|
|
try:
|
|
def content_generator():
|
|
for chunk in stream_response:
|
|
if should_stop():
|
|
stream_response.response.close()
|
|
break
|
|
if content := chunk.data.choices[0].delta.content:
|
|
yield content
|
|
yield from skip_prefix(content_generator(), continuation_text)
|
|
finally:
|
|
stream_response.response.close()
|
|
|
|
def token_count(self, system_prompt: str, main_context: str) -> int:
|
|
messages = [
|
|
SystemMessage(content=system_prompt),
|
|
UserMessage(content=main_context),
|
|
]
|
|
|
|
tokenized = self._tokenizer.encode_chat_completion(
|
|
ChatCompletionRequest(
|
|
messages=messages,
|
|
model=self._model
|
|
)
|
|
)
|
|
return len(tokenized.tokens)
|
|
|
|
def token_limit(self) -> int:
|
|
return self._token_limit
|