70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
from typing import Iterator, Optional
|
|
from abc import ABC, abstractmethod
|
|
import os
|
|
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 .llm_engine import LlmEngine
|
|
|
|
class MistralLlmEngine(LlmEngine):
|
|
def __init__(
|
|
self,
|
|
model: str,
|
|
temperature: float,
|
|
api_key: str,
|
|
token_limit: int
|
|
):
|
|
self._model = model
|
|
self._temperature = temperature
|
|
self._api_key = api_key
|
|
self._token_limit = token_limit
|
|
self._client = Mistral(api_key=api_key)
|
|
self._tokenizer = MistralTokenizer.v3()
|
|
|
|
def infer(self, system_prompt: str, main_context: str) -> Iterator[str]:
|
|
messages = [
|
|
SystemMessage(content=system_prompt),
|
|
UserMessage(content=main_context),
|
|
]
|
|
messages = [
|
|
{
|
|
"role": "system",
|
|
"content": system_prompt,
|
|
},
|
|
{
|
|
"role": "user",
|
|
"content": main_context,
|
|
},
|
|
{
|
|
"role": "assistant",
|
|
"content": "<",
|
|
"prefix": True,
|
|
},
|
|
]
|
|
stream_response = self._client.chat.stream(
|
|
model=self._model,
|
|
messages=messages,
|
|
temperature=self._temperature,
|
|
)
|
|
for chunk in stream_response:
|
|
if content := chunk.data.choices[0].delta.content:
|
|
yield chunk.data.choices[0].delta.content
|
|
|
|
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 |