made tokenizer concurrent and threadsafe

app/MindWork AI Studio/Components/ChatComponent.razor.cs

@@ -1259,12 +1259,12 @@ public partial class ChatComponent : MSGComponentBase, IAsyncDisposable
             return;
         }
 
-        var response = await this.RustService.GetTokenCount(this.Provider.InstanceName, this.Provider.TokenizerPath, this.inputField.Value);
+        var response = await this.RustService.GetTokenCount(this.Provider, this.inputField.Value);
         if (response is null)
             return;
         if (!response.Value.Success)
         {
-            this.Logger.LogWarning($"Failed to calculate token count: status='{response.Value.Status}', reason='{response.Value.Message}'");
+            this.Logger.LogWarning("Failed to calculate token count: reason='{Reason}'", response.Value.Message);
             return;
         }
         this.tokenCount = response.Value.TokenCount.ToString();

app/MindWork AI Studio/Tools/Rust/TokenizerResponse.cs

@@ -4,5 +4,4 @@ public readonly record struct TokenizerResponse(
     bool Success,
     int TokenCount,
     string Message,
-    TokenizerStatus Status = TokenizerStatus.UNAVAILABLE,
     string StoredPath = "");

app/MindWork AI Studio/Tools/Rust/TokenizerStatus.cs

@@ -1,8 +0,0 @@
-namespace AIStudio.Tools.Rust;
-
-public enum TokenizerStatus
-{
-    UNAVAILABLE,
-    RUNNING,
-    AVAILABLE,
-}

app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs

@@ -66,7 +66,7 @@ public sealed partial class DataSourceEmbeddingService
     {
         var segments = new List<ExtractedFileSegment>();
 
-        await foreach (var segment in rustService.StreamArbitraryFileDataWithTokenCounts(filePath, embeddingProvider.Name, embeddingProvider.TokenizerPath, token))
+        await foreach (var segment in rustService.StreamArbitraryFileDataWithTokenCounts(filePath, embeddingProvider, token))
         {
             var normalized = NormalizeChunkSegment(segment.Content);
             if (!string.IsNullOrWhiteSpace(normalized))
@@ -536,8 +536,8 @@ public sealed partial class DataSourceEmbeddingService
 
     private async Task<int> GetEmbeddingTokenCountAsync(EmbeddingProvider embeddingProvider, string text, CancellationToken token)
     {
-        var response = await rustService.GetTokenCount(embeddingProvider.Name, embeddingProvider.TokenizerPath, text, token);
+        var response = await rustService.GetTokenCount(embeddingProvider, text, token);
-        if (response is { Success: true, Status: TokenizerStatus.AVAILABLE })
+        if (response is { Success: true })
             return response.Value.TokenCount;
 
         var message = response?.Message ?? "No response was returned by the tokenizer service.";

app/MindWork AI Studio/Tools/Services/DataSourceLocalRetrievalService.cs

@@ -6,7 +6,6 @@ using AIStudio.Tools.Databases;
 using AIStudio.Tools.Databases.EmbeddingState;
 using AIStudio.Tools.Databases.VectorStore;
 using AIStudio.Tools.RAG;
-using AIStudio.Tools.Rust;
 
 namespace AIStudio.Tools.Services;
 
@@ -160,12 +159,8 @@ public sealed class DataSourceLocalRetrievalService(
             return false;
         }
 
-        var tokenCountResponse = await rustService.GetTokenCount(
+        var tokenCountResponse = await rustService.GetTokenCount(embeddingProvider, query, token);
-            embeddingProvider.Name,
+        if (tokenCountResponse is not { Success: true })
-            embeddingProvider.TokenizerPath,
-            query,
-            token);
-        if (tokenCountResponse is not { Success: true, Status: TokenizerStatus.AVAILABLE })
         {
             logger.LogWarning(
                 "Skipping vector retrieval for data source '{DataSourceName}' ({DataSourceId}) because the token count for embedding provider '{EmbeddingProviderName}' could not be determined. Reason='{Reason}'.",

app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs

@@ -2,6 +2,7 @@ using System.Text;
 using System.Text.Json;
 using System.Runtime.CompilerServices;
 
+using AIStudio.Settings;
 using AIStudio.Tools.Rust;
 
 namespace AIStudio.Tools.Services;
@@ -85,37 +86,21 @@ public sealed partial class RustService
 
     public async IAsyncEnumerable<string> StreamArbitraryFileData(string path, bool extractImages = false, [EnumeratorCancellation] CancellationToken token = default)
     {
-        await foreach (var segment in this.StreamArbitraryFileDataCore(path, extractImages, false, token))
+        await foreach (var segment in this.StreamArbitraryFileDataCore(path, extractImages, false, string.Empty, token))
             yield return segment.Content;
     }
 
     public async IAsyncEnumerable<ArbitraryFileDataSegment> StreamArbitraryFileDataWithTokenCounts(
         string path,
-        string providerName,
+        EmbeddingProvider embeddingProvider,
-        string tokenizerPath,
         [EnumeratorCancellation] CancellationToken token = default)
     {
-        await this.tokenizerLock.WaitAsync(token);
+        await foreach (var segment in this.StreamArbitraryFileDataCore(path, false, true, embeddingProvider.TokenizerPath, token))
-        try
         {
-            var tokenizerResponse = await this.EnsureTokenizerCoreAsync(providerName, tokenizerPath);
+            if (segment.TokenCount is null)
-            if (tokenizerResponse is not { Success: true, Status: TokenizerStatus.AVAILABLE })
+                throw new InvalidOperationException($"Rust did not return a token count for an extracted segment from '{path}' using provider '{embeddingProvider.Name}'.");
-            {
-                var message = tokenizerResponse?.Message ?? "No response was returned by the tokenizer service.";
-                throw new InvalidOperationException($"Could not initialize tokenizer for provider '{providerName}'. {message}");
-            }
 
-            await foreach (var segment in this.StreamArbitraryFileDataCore(path, false, true, token))
+            yield return new(segment.Content, segment.TokenCount.Value);
-            {
-                if (segment.TokenCount is null)
-                    throw new InvalidOperationException($"Rust did not return a token count for an extracted segment from '{path}'.");
-
-                yield return new(segment.Content, segment.TokenCount.Value);
-            }
-        }
-        finally
-        {
-            this.tokenizerLock.Release();
         }
     }
 
@@ -123,10 +108,11 @@ public sealed partial class RustService
         string path,
         bool extractImages,
         bool includeTokenCount,
+        string tokenizerPath,
         [EnumeratorCancellation] CancellationToken token)
     {
         var streamId = Guid.NewGuid().ToString();
-        var requestUri = $"/retrieval/fs/extract?path={Uri.EscapeDataString(path)}&stream_id={streamId}&extract_images={extractImages}&include_token_count={includeTokenCount}";
+        var requestUri = $"/retrieval/fs/extract?path={Uri.EscapeDataString(path)}&stream_id={streamId}&extract_images={extractImages}&include_token_count={includeTokenCount}&tokenizer_path={Uri.EscapeDataString(tokenizerPath)}";
         using var request = new HttpRequestMessage(HttpMethod.Get, requestUri);
         using var response = await this.http.SendAsync(request, HttpCompletionOption.ResponseHeadersRead, token);
 

app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs

@@ -1,4 +1,4 @@
-using AIStudio.Provider;
+using AIStudio.Settings;
 using AIStudio.Tools.Rust;
 
 namespace AIStudio.Tools.Services;
@@ -7,35 +7,12 @@ public sealed partial class RustService
 {
     internal const int MAX_TOKEN_COUNT_REQUEST_TEXT_LENGTH = 200_000;
 
-    private readonly SemaphoreSlim tokenizerLock = new(1, 1);
-    private string currentTokenizerPath = string.Empty;
-    private bool hasInitializedTokenizer;
-
     private static TokenizerResponse CreateUnavailableTokenizerResponse(string message) => new(
         false,
         0,
         message,
-        TokenizerStatus.UNAVAILABLE,
         string.Empty);
 
-    public async Task<TokenizerResponse> GetTokenizerInfo(CancellationToken cancellationToken = default)
-    {
-        try
-        {
-            return await this.http.GetFromJsonAsync<TokenizerResponse>("/system/tokenizer/info", this.jsonRustSerializerOptions, cancellationToken);
-        }
-        catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
-        {
-            this.logger?.LogWarning("Fetching tokenizer info from Rust service was cancelled by caller.");
-            return CreateUnavailableTokenizerResponse("Operation cancelled by caller.");
-        }
-        catch (Exception e)
-        {
-            this.logger?.LogError(e, "Error while fetching tokenizer info from Rust service.");
-            return CreateUnavailableTokenizerResponse(e.Message);
-        }
-    }
-
     public async Task<TokenizerResponse> ValidateTokenizer(string filePath)
     {
         var result = await this.http.PostAsJsonAsync("/tokenizer/validate", new {
@@ -86,99 +63,25 @@ public sealed partial class RustService
         return await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions);
     }
     
-    public Task<TokenizerResponse?> GetTokenCount(string text)
+    public Task<TokenizerResponse?> GetTokenCount(AIStudio.Settings.Provider provider, string text, CancellationToken cancellationToken = default) =>
-    {
+        this.GetTokenCount(provider.InstanceName, provider.TokenizerPath, text, cancellationToken);
-        return this.GetTokenCountCoreAsync(text);
-    }
 
-    public async Task<TokenizerResponse?> GetTokenCount(string providerName, string path, string text, CancellationToken cancellationToken = default)
+    public Task<TokenizerResponse?> GetTokenCount(EmbeddingProvider provider, string text, CancellationToken cancellationToken = default) =>
-    {
+        this.GetTokenCount(provider.Name, provider.TokenizerPath, text, cancellationToken);
-        await this.tokenizerLock.WaitAsync(cancellationToken);
-        try
-        {
-            var tokenizerResponse = await this.EnsureTokenizerCoreAsync(providerName, path);
-            if (tokenizerResponse is not { Success: true, Status: TokenizerStatus.AVAILABLE })
-                return tokenizerResponse;
 
-            return await this.GetTokenCountCoreAsync(text, cancellationToken);
+    private async Task<TokenizerResponse?> GetTokenCount(string providerName, string tokenizerPath, string text, CancellationToken cancellationToken)
-        }
-        finally
-        {
-            this.tokenizerLock.Release();
-        }
-    }
-
-    private async Task<TokenizerResponse?> GetTokenCountCoreAsync(string text, CancellationToken cancellationToken = default)
     {
         var result = await this.http.PostAsJsonAsync("/tokenizer/count", new {
             text = text,
+            tokenizer_path = tokenizerPath,
         }, this.jsonRustSerializerOptions, cancellationToken);
 
         if (!result.IsSuccessStatusCode)
         {
-            this.logger!.LogError($"Failed to get the token count '{result.StatusCode}'");
+            this.logger!.LogError("Failed to get the token count for provider '{ProviderName}': {StatusCode}", providerName, result.StatusCode);
-            this.hasInitializedTokenizer = false;
             return CreateUnavailableTokenizerResponse("Error while getting token count from Rust service: "+result.StatusCode);
         }
 
-        var response = await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions);
+        return await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions, cancellationToken);
-        if (response is not { Status: TokenizerStatus.AVAILABLE })
-            this.hasInitializedTokenizer = false;
-
-        return response;
-    }
-
-    public async Task<TokenizerResponse?> SetTokenizer(string providerName, string path)
-    {
-        this.logger!.LogInformation($"Setting a new tokenizer for '{providerName}'");
-        var result = await this.http.PostAsJsonAsync("/tokenizer/set", new {
-            file_path = path,
-        }, this.jsonRustSerializerOptions);
-
-        if (!result.IsSuccessStatusCode)
-        {
-            this.logger!.LogError($"Failed to set the tokenizer '{result.StatusCode}'");
-            this.hasInitializedTokenizer = false;
-            return CreateUnavailableTokenizerResponse("An error occured while sending the path to the Rust framework for setting a tokenizer: "+result.StatusCode);
-        }
-
-        var response = await result.Content.ReadFromJsonAsync<TokenizerResponse>(this.jsonRustSerializerOptions);
-        if (response is not { Success: true, Status: TokenizerStatus.AVAILABLE })
-            this.hasInitializedTokenizer = false;
-
-        return response;
-    }
-
-    public async Task<TokenizerResponse?> EnsureTokenizer(string providerName, string path)
-    {
-        await this.tokenizerLock.WaitAsync();
-        try
-        {
-            return await this.EnsureTokenizerCoreAsync(providerName, path);
-        }
-        finally
-        {
-            this.tokenizerLock.Release();
-        }
-    }
-
-    private async Task<TokenizerResponse?> EnsureTokenizerCoreAsync(string providerName, string path)
-    {
-        if (this.hasInitializedTokenizer && this.currentTokenizerPath == path)
-            return new TokenizerResponse(true, 0, string.Empty, TokenizerStatus.AVAILABLE);
-
-        var response = await this.SetTokenizer(providerName, path);
-        if (response is { Success: true, Status: TokenizerStatus.AVAILABLE })
-        {
-            this.currentTokenizerPath = path;
-            this.hasInitializedTokenizer = true;
-        }
-        else
-        {
-            this.currentTokenizerPath = string.Empty;
-            this.hasInitializedTokenizer = false;
-        }
-
-        return response;
     }
 }

app/MindWork AI Studio/Tools/Services/RustService.cs

@@ -94,7 +94,6 @@ public sealed partial class RustService : BackgroundService
     {
         this.http.Dispose();
         this.userLanguageLock.Dispose();
-        this.tokenizerLock.Dispose();
         this.userNameLock.Dispose();
         base.Dispose();
     }

runtime/src/file_data.rs

@@ -22,6 +22,7 @@ use log::{debug, error, warn};
 use tokio::io::AsyncBufReadExt;
 use tokio::sync::mpsc;
 use tokio_stream::wrappers::ReceiverStream;
+use tokenizers::tokenizer::Tokenizer;
 
 #[derive(Debug, Serialize)]
 pub struct Chunk {
@@ -39,8 +40,8 @@ impl Chunk {
     
     pub fn set_stream_id(&mut self, stream_id: &str) { self.stream_id = stream_id.to_string(); }
 
-    pub fn set_token_count(&mut self) -> std::result::Result<(), String> {
+    pub fn set_token_count(&mut self, tokenizer: &Tokenizer) -> std::result::Result<(), String> {
-        self.token_count = Some(crate::tokenizer::get_segment_token_count(&self.content)?);
+        self.token_count = Some(crate::tokenizer::get_segment_token_count(tokenizer, &self.content)?);
         Ok(())
     }
 
@@ -116,6 +117,8 @@ pub struct ExtractDataQuery {
     extract_images: bool,
     #[serde(default, deserialize_with = "deserialize_bool_case_insensitive")]
     include_token_count: bool,
+    #[serde(default)]
+    tokenizer_path: String,
 }
 
 fn deserialize_bool_case_insensitive<'de, D>(deserializer: D) -> std::result::Result<bool, D::Error>
@@ -165,7 +168,18 @@ pub async fn extract_data(
 
     let stream = stream! {
         match query {
-            Ok(query) => {
+            Ok(query) => 'request: {
+                let tokenizer = if query.include_token_count {
+                    match crate::tokenizer::get_tokenizer(&query.tokenizer_path) {
+                        Ok(tokenizer) => Some(tokenizer),
+                        Err(e) => {
+                            yield Ok(Event::default().json_data(format!("Error loading tokenizer: {e}")).unwrap_or_else(|_| Event::default().data(format!("Error loading tokenizer: {e}"))));
+                            break 'request;
+                        },
+                    }
+                } else {
+                    None
+                };
                 let stream_result = stream_data(&query.path, query.extract_images).await;
                 let id_ref = &query.stream_id;
 
@@ -178,8 +192,8 @@ pub async fn extract_data(
 
                                     for mut chunk in chunks {
                                         chunk.set_stream_id(id_ref);
-                                        if query.include_token_count {
+                                        if let Some(tokenizer) = tokenizer.as_deref() {
-                                            if let Err(e) = chunk.set_token_count() {
+                                            if let Err(e) = chunk.set_token_count(tokenizer) {
                                                 yield Ok(Event::default().json_data(format!("Error counting tokens: {e}")).unwrap_or_else(|_| Event::default().data(format!("Error counting tokens: {e}"))));
                                                 break 'stream_chunks;
                                             }

runtime/src/runtime_api.rs

@@ -32,7 +32,6 @@ pub fn start_runtime_api() {
     let app = Router::new()
         .route("/system/dotnet/port", get(crate::dotnet::dotnet_port))
         .route("/system/dotnet/ready", get(crate::dotnet::dotnet_ready))
-        .route("/system/tokenizer/info", get(crate::tokenizer::tokenizer_info))
         .route("/system/qdrant-edge/info", get(crate::qdrant_edge_database::qdrant_edge_info))
         .route("/system/qdrant-edge/ensure", post(crate::qdrant_edge_database::ensure_qdrant_edge_store))
         .route("/system/qdrant-edge/insert", post(crate::qdrant_edge_database::insert_qdrant_edge_embedding))
@@ -73,7 +72,6 @@ pub fn start_runtime_api() {
         .route("/tokenizer/validate", post(crate::tokenizer::validate_tokenizer))
         .route("/tokenizer/store", post(crate::tokenizer::store_tokenizer))
         .route("/tokenizer/delete", post(crate::tokenizer::delete_tokenizer))
-        .route("/tokenizer/set", post(crate::tokenizer::set_tokenizer))
         .route("/shortcuts/register", post(crate::app_window::register_shortcut))
         .route("/shortcuts/validate", post(crate::app_window::validate_shortcut))
         .route("/shortcuts/suspend", post(crate::app_window::suspend_shortcuts))

runtime/src/tokenizer.rs

@@ -1,10 +1,10 @@
+use std::collections::HashMap;
 use std::fs;
 use std::path::PathBuf;
-use std::sync::{Mutex, OnceLock, RwLock};
+use std::sync::{Arc, Mutex, OnceLock, RwLock};
 
 use axum::Json;
 use log::{error, warn};
-use once_cell::sync::Lazy;
 use serde::{Deserialize, Serialize};
 use tauri::path::BaseDirectory;
 use tauri::Manager;
@@ -14,30 +14,15 @@ use crate::api_token::APIToken;
 use crate::environment::DATA_DIRECTORY;
 
 const DEFAULT_TOKENIZER_RESOURCE_PATH: &str = "resources/tokenizers/tokenizer.json";
-const NO_TOKENIZER_LOADED_MESSAGE: &str = "Tokenizer must be set before counting tokens.";
 
-static TOKENIZER: OnceLock<RwLock<Option<Tokenizer>>> = OnceLock::new();
+static TOKENIZERS: OnceLock<RwLock<HashMap<PathBuf, Arc<Tokenizer>>>> = OnceLock::new();
 static DEFAULT_TOKENIZER_PATH: OnceLock<PathBuf> = OnceLock::new();
-static TOKENIZER_STATUS: Lazy<Mutex<TokenizerStatusInfo>> = Lazy::new(|| Mutex::new(TokenizerStatusInfo::default()));
+static TOKENIZER_STORAGE_LOCK: Mutex<()> = Mutex::new(());
-static TOKENIZER_OPERATION_LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
-
-#[derive(Clone, Copy, Default, Serialize, PartialEq, Eq)]
-pub enum TokenizerStatus {
-    #[default]
-    Unavailable,
-    Running,
-    Available,
-}
-
-#[derive(Default)]
-struct TokenizerStatusInfo {
-    status: TokenizerStatus,
-    unavailable_reason: Option<String>,
-}
 
 #[derive(Deserialize)]
 pub struct SetTokenText {
     text: String,
+    tokenizer_path: String,
 }
 
 #[derive(Clone, Deserialize)]
@@ -61,7 +46,6 @@ pub struct TokenizerResponse {
     success: bool,
     token_count: usize,
     message: String,
-    status: TokenizerStatus,
     stored_path: String,
 }
 
@@ -71,7 +55,6 @@ impl TokenizerResponse {
             success: true,
             token_count,
             message: String::new(),
-            status: TokenizerStatus::Available,
             stored_path: String::new(),
         }
     }
@@ -81,7 +64,6 @@ impl TokenizerResponse {
             success: true,
             token_count: 0,
             message: String::new(),
-            status: TokenizerStatus::Available,
             stored_path,
         }
     }
@@ -91,7 +73,6 @@ impl TokenizerResponse {
             success: false,
             token_count: 0,
             message: reason,
-            status: TokenizerStatus::Unavailable,
             stored_path: String::new(),
         }
     }
@@ -106,7 +87,6 @@ pub fn set_default_tokenizer_path(app_handle: tauri::AppHandle) {
         Err(e) => {
             let reason = format!("The default tokenizer file '{DEFAULT_TOKENIZER_RESOURCE_PATH}' could not be resolved: {e}");
             error!(Source = "Tokenizer"; "{reason}");
-            set_tokenizer_unavailable(reason);
             return;
         }
     };
@@ -114,7 +94,6 @@ pub fn set_default_tokenizer_path(app_handle: tauri::AppHandle) {
     if !tokenizer_path.is_file() {
         let reason = format!("The default tokenizer file was not found: {}", tokenizer_path.display());
         error!(Source = "Tokenizer"; "{reason}");
-        set_tokenizer_unavailable(reason);
         return;
     }
 
@@ -124,23 +103,8 @@ pub fn set_default_tokenizer_path(app_handle: tauri::AppHandle) {
     }
 }
 
-pub async fn tokenizer_info(_token: APIToken) -> Json<TokenizerResponse> {
-    let status = TOKENIZER_STATUS.lock().unwrap();
-    match status.status {
-        TokenizerStatus::Available => Json(TokenizerResponse::available(0)),
-        TokenizerStatus::Running => Json(TokenizerResponse {
-            success: false,
-            token_count: 0,
-            message: String::new(),
-            status: TokenizerStatus::Running,
-            stored_path: String::new(),
-        }),
-        TokenizerStatus::Unavailable => Json(TokenizerResponse::unavailable(status.unavailable_reason.clone().unwrap_or_default())),
-    }
-}
-
 pub async fn token_count(_token: APIToken, req: Json<SetTokenText>) -> Json<TokenizerResponse> {
-    match get_token_count(&req.text) {
+    match get_token_count(&req.tokenizer_path, &req.text) {
         Ok(count) => Json(TokenizerResponse::available(count)),
         Err(e) => Json(TokenizerResponse::unavailable(e)),
     }
@@ -167,92 +131,29 @@ pub async fn delete_tokenizer(_token: APIToken, payload: Json<TokenizerDelete>)
     }
 }
 
-pub async fn set_tokenizer(_token: APIToken, payload: Json<TokenizerPath>) -> Json<TokenizerResponse> {
-    match handle_tokenizer_set(&payload.file_path) {
-        Ok(_) => Json(TokenizerResponse::available(0)),
-        Err(e) => Json(TokenizerResponse::unavailable(e)),
-    }
-}
-
-pub fn handle_tokenizer_set(path: &str) -> Result<(), String> {
-    let _operation_guard = begin_tokenizer_operation()?;
-    set_tokenizer_running();
-
-    let tokenizer_path = resolve_tokenizer_path(path).map_err(|e| {
-        error!(Source = "Tokenizer"; "{e} Starting the app without a tokenizer.");
-        unavailable_with_status_update(&e)
-    })?;
-
-    let tokenizer = load_tokenizer_from_file(&tokenizer_path).map_err(|e| {
-        error!(Source = "Tokenizer"; "{e}");
-        unavailable_with_status_update(&e)
-    })?;
-
-    match tokenizer_state().write() {
-        Ok(mut tokenizer_guard) => *tokenizer_guard = Some(tokenizer),
-        Err(_) => return Err(unavailable_with_status_update("Tokenizer state lock is poisoned.")),
-    }
-
-    set_tokenizer_available();
-    Ok(())
-}
-
 fn handle_tokenizer_validate(path: &PathBuf) -> Result<usize, String> {
-    let _operation_guard = begin_tokenizer_operation()?;
+    validate_tokenizer_file(path)
-    set_tokenizer_running();
-
-    let result = validate_tokenizer_file(path);
-    match tokenizer_state().read() {
-        Ok(tokenizer_guard) if tokenizer_guard.is_some() => set_tokenizer_available(),
-        Ok(_) => set_tokenizer_unavailable(NO_TOKENIZER_LOADED_MESSAGE.to_string()),
-        Err(_) => set_tokenizer_unavailable("Tokenizer state lock is poisoned.".to_string()),
-    }
-
-    result
 }
 
-pub fn get_token_count(text: &str) -> Result<usize, String> {
+pub fn get_token_count(path: &str, text: &str) -> Result<usize, String> {
-    get_token_count_internal(text, true)
+    let tokenizer = get_tokenizer(path)?;
+    get_token_count_internal(&tokenizer, text, true)
 }
 
-pub fn get_segment_token_count(text: &str) -> Result<usize, String> {
+pub fn get_segment_token_count(tokenizer: &Tokenizer, text: &str) -> Result<usize, String> {
     // Special tokens belong to the final encoding and would inflate sums across many segments.
-    get_token_count_internal(text, false)
+    get_token_count_internal(tokenizer, text, false)
 }
 
-fn get_token_count_internal(text: &str, add_special_tokens: bool) -> Result<usize, String> {
+fn get_token_count_internal(tokenizer: &Tokenizer, text: &str, add_special_tokens: bool) -> Result<usize, String> {
     if text.trim().is_empty() {
         return Ok(0);
     }
 
-    let _operation_guard = begin_tokenizer_operation()?;
+    tokenizer
-    {
+        .encode(text, add_special_tokens)
-        let status = TOKENIZER_STATUS.lock().unwrap();
+        .map(|encoding| encoding.len())
-        if status.status != TokenizerStatus::Available {
+        .map_err(|e| format!("Failed to tokenize text: {e}"))
-            return Err(status.unavailable_reason.clone().unwrap_or_else(|| NO_TOKENIZER_LOADED_MESSAGE.to_string()));
-        }
-    }
-
-    let tokenizer_guard = tokenizer_state()
-        .read()
-        .map_err(|_| unavailable_with_status_update("Tokenizer state lock is poisoned."))?;
-    let tokenizer = match tokenizer_guard.as_ref() {
-        Some(tokenizer) => tokenizer,
-        None => {
-            drop(tokenizer_guard);
-            return Err(unavailable_with_status_update("Tokenizer not initialized."));
-        }
-    };
-    let token_count = match tokenizer.encode(text, add_special_tokens) {
-        Ok(enc) => enc.len(),
-        Err(e) => {
-            let reason = format!("Failed to tokenize text: {e}");
-            drop(tokenizer_guard);
-            return Err(unavailable_with_status_update(&reason));
-        }
-    };
-
-    Ok(token_count)
 }
 
 fn validate_tokenizer_file(path: &PathBuf) -> Result<usize, String> {
@@ -293,7 +194,11 @@ fn handle_tokenizer_store(payload: &TokenizerStorage) -> Result<String, std::io:
         return Ok(destination_path.to_string_lossy().to_string());
     }
 
+    let _storage_guard = TOKENIZER_STORAGE_LOCK
+        .lock()
+        .map_err(|_| std::io::Error::other("Tokenizer storage lock is poisoned."))?;
     if model_path.try_exists()? {
+        invalidate_tokenizers_under(&model_path);
         fs::remove_dir_all(&model_path)?;
     }
 
@@ -320,50 +225,63 @@ fn handle_tokenizer_delete(payload: &TokenizerDelete) -> Result<(), std::io::Err
         .join("tokenizers")
         .join(&payload.model_id);
 
+    let _storage_guard = TOKENIZER_STORAGE_LOCK
+        .lock()
+        .map_err(|_| std::io::Error::other("Tokenizer storage lock is poisoned."))?;
     if tokenizer_path.exists() {
+        invalidate_tokenizers_under(&tokenizer_path);
         fs::remove_dir_all(tokenizer_path)?;
     }
 
     Ok(())
 }
 
-fn tokenizer_state() -> &'static RwLock<Option<Tokenizer>> {
+fn tokenizer_cache() -> &'static RwLock<HashMap<PathBuf, Arc<Tokenizer>>> {
-    TOKENIZER.get_or_init(|| RwLock::new(None))
+    TOKENIZERS.get_or_init(|| RwLock::new(HashMap::new()))
 }
 
-fn begin_tokenizer_operation() -> Result<std::sync::MutexGuard<'static, ()>, String> {
+pub fn get_tokenizer(path: &str) -> Result<Arc<Tokenizer>, String> {
-    TOKENIZER_OPERATION_LOCK
+    let resolved_path = resolve_tokenizer_path(path)?;
-        .lock()
+    let tokenizer_path = fs::canonicalize(&resolved_path)
-        .map_err(|_| unavailable_with_status_update("Tokenizer operation lock is poisoned."))
+        .map_err(|e| format!("Could not resolve tokenizer file '{}': {e}", resolved_path.display()))?;
-}
 
-fn set_tokenizer_available() {
+    if let Some(tokenizer) = tokenizer_cache()
-    let mut status = TOKENIZER_STATUS.lock().unwrap();
+        .read()
-    status.status = TokenizerStatus::Available;
+        .map_err(|_| "Tokenizer cache lock is poisoned.".to_string())?
-    status.unavailable_reason = None;
+        .get(&tokenizer_path)
-}
+        .cloned()
-
+    {
-fn set_tokenizer_running() {
+        return Ok(tokenizer);
-    let mut status = TOKENIZER_STATUS.lock().unwrap();
-    status.status = TokenizerStatus::Running;
-    status.unavailable_reason = None;
-}
-
-fn set_tokenizer_unavailable(reason: String) {
-    let mut status = TOKENIZER_STATUS.lock().unwrap();
-    status.status = TokenizerStatus::Unavailable;
-    status.unavailable_reason = Some(reason);
-}
-
-fn unavailable_with_status_update(reason: &str) -> String {
-    let reason = reason.to_string();
-    match tokenizer_state().write() {
-        Ok(mut tokenizer_guard) => *tokenizer_guard = None,
-        Err(_) => set_tokenizer_unavailable("Tokenizer state lock is poisoned.".to_string()),
     }
 
-    set_tokenizer_unavailable(reason.clone());
+    let _storage_guard = TOKENIZER_STORAGE_LOCK
-    reason
+        .lock()
+        .map_err(|_| "Tokenizer storage lock is poisoned.".to_string())?;
+    if let Some(tokenizer) = tokenizer_cache()
+        .read()
+        .map_err(|_| "Tokenizer cache lock is poisoned.".to_string())?
+        .get(&tokenizer_path)
+        .cloned()
+    {
+        return Ok(tokenizer);
+    }
+
+    let loaded_tokenizer = Arc::new(load_tokenizer_from_file(&tokenizer_path)?);
+    let mut cache = tokenizer_cache()
+        .write()
+        .map_err(|_| "Tokenizer cache lock is poisoned.".to_string())?;
+    Ok(cache
+        .entry(tokenizer_path)
+        .or_insert_with(|| loaded_tokenizer)
+        .clone())
+}
+
+fn invalidate_tokenizers_under(path: &PathBuf) {
+    let cache_path = fs::canonicalize(path).unwrap_or_else(|_| path.clone());
+    match tokenizer_cache().write() {
+        Ok(mut cache) => cache.retain(|tokenizer_path, _| !tokenizer_path.starts_with(&cache_path)),
+        Err(_) => warn!(Source = "Tokenizer"; "Could not invalidate tokenizer cache because its lock is poisoned."),
+    }
 }
 
 fn resolve_tokenizer_path(path: &str) -> Result<PathBuf, String> {