Added mistral llm engine

install.sh

@@ -20,4 +20,4 @@ vim .env
 cd /root/SIA
 python3 -m sia
 
-#The SIA source is located in /SIA. Not all features are implemented yet. Look at the readme and code to find what is missing. Make sure to unit test your work.
+#The SIA source is located in /root/sia. Not all features are implemented yet. Look at the readme and code to find what is missing. Make sure to unit test your work.

requirements.txt

@@ -1,6 +1,8 @@
 accelerate
 aiohttp
 bs4
+mistral_common
+mistralai
 openai
 python-dotenv
 tiktoken

sia/__main__.py

@@ -9,6 +9,7 @@ from .config import Config
 from .hf_llm_engine import HfLlmEngine
 from .llm_engine import LlmEngine
 from .local_llm_engine import LocalLlmEngine
+from .mistral_llm_engine import MistralLlmEngine
 from .openai_llm_engine import OpenAILlmEngine
 from .response_parser import ResponseParser
 from .system_metrics import SystemMetrics
@@ -39,8 +40,8 @@ class Main:
                 self._llm = LocalLlmEngine(
                     self._config.model,
                     self._config.temperature,
+                    self._config.token_limit,
                     self._config.api_token,
-                    self._config.token_limit
                 )
             case "hf":
                 self._llm = HfLlmEngine(
@@ -50,6 +51,13 @@ class Main:
                 )
             case "openai":
                 self._llm = OpenAILlmEngine(
+                    self._config.model,
+                    self._config.temperature,
+                    self._config.token_limit,
+                    self._config.api_token,
+                )
+            case "mistral":
+                self._llm = MistralLlmEngine(
                     self._config.model,
                     self._config.temperature,
                     self._config.api_token,

sia/local_llm_engine.py

@@ -26,7 +26,7 @@ class LocalLlmEngine(LlmEngine):
         """
         self._temperature = temperature
         self._token_limit = token_limit
-        self._tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self._tokenizer = AutoTokenizer.from_pretrained(model_path, token=api_token)
         model = AutoModelForCausalLM.from_pretrained(
             model_path,
             return_dict=True,

sia/mistral_llm_engine.py

@@ -0,0 +1,70 @@
+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

sia/openai_llm_engine.py

@@ -14,8 +14,8 @@ class OpenAILlmEngine(LlmEngine):
         self,
         model: str,
         temperature: float,
+        token_limit: int,
         api_key: str,
-        token_limit: int = 0,
     ):
         """
         Initialize the OpenAI LLM Engine.