updadted mistral token counter

app.py

@@ -3,7 +3,6 @@ import os
 import tiktoken
 import streamlit as st
 
-# from llama_index.llms.gemini import Gemini
 from llama_index.core import (
     VectorStoreIndex,
     Settings,
@@ -14,10 +13,33 @@ from streamlit_pdf_viewer import pdf_viewer
 
 MAX_OUTPUT_TOKENS = 2048  
 
+class MistralTokens:
+    """
+    Returns tokens for MistralAI models.
+    
+    See: https://docs.mistral.ai/guides/tokenization/
+    """
+    def __init__(self, llm_name):
+        from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+        self.tokenizer = MistralTokenizer.from_model(llm_name)
 
-class CountGeminiTokens:
+    def __call__(self, input):
+        """This returns all the tokens indices in a list since LlamaIndex seems to count by calling `len()` on the tokenizer function."""
+        from mistral_common.protocol.instruct.messages import UserMessage
+        from mistral_common.protocol.instruct.request import ChatCompletionRequest
+
+        return self.tokenizer.encode_chat_completion(
+            ChatCompletionRequest(
+                tools=[], 
+                messages=[
+                    UserMessage(content=input)
+                ]
+            )
+        ).tokens
+
+class GeminiTokens:
     """
-    Count tokens in Gemini models.
+    Returns tokens for Gemini models.
     
     See: https://medium.com/google-cloud/counting-gemini-text-tokens-locally-with-the-vertex-ai-sdk-78979fea6244
     """
@@ -99,7 +121,6 @@ def main():
         # https://docs.llamaindex.ai/en/stable/module_guides/models/llms/
         if llm_key is not None:
             if provider == 'google':
-                # raise NotImplementedError(f"{provider} is not supported yet")
                 from llama_index.llms.gemini import Gemini
                 from llama_index.embeddings.gemini import GeminiEmbedding
 
@@ -110,7 +131,7 @@ def main():
                     temperature=temperature,
                     max_tokens=MAX_OUTPUT_TOKENS
                 )
-                Settings.tokenizer = CountGeminiTokens(llm_name) #tokenization.get_tokenizer_for_model(llm_name).compute_tokens
+                Settings.tokenizer = GeminiTokens(llm_name) 
                 Settings.num_output = MAX_OUTPUT_TOKENS
                 Settings.embed_model = GeminiEmbedding(
                     model_name="models/text-embedding-004", api_key=os.environ.get("GOOGLE_API_KEY") #, title="this is a document"
@@ -141,7 +162,7 @@ def main():
             elif provider == 'mistralai':
                 from llama_index.llms.mistralai import MistralAI
                 from llama_index.embeddings.mistralai import MistralAIEmbedding
-
+                
                 os.environ['MISTRAL_API_KEY'] = str(llm_key)
                 Settings.llm = MistralAI(
                     model=llm_name, 
@@ -150,7 +171,7 @@ def main():
                     random_seed=42,
                     safe_mode=True
                 )
-                # Settings.tokenizer = tiktoken.encoding_for_model(llm_name).encode
+                Settings.tokenizer = MistralTokens(llm_name)
                 Settings.num_output = MAX_OUTPUT_TOKENS
                 Settings.embed_model = MistralAIEmbedding(
                     model_name="mistral-embed",