diff --git a/app/MindWork AI Studio/Components/ChatComponent.razor.cs b/app/MindWork AI Studio/Components/ChatComponent.razor.cs index 46edd698..5f7fabff 100644 --- a/app/MindWork AI Studio/Components/ChatComponent.razor.cs +++ b/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(); diff --git a/app/MindWork AI Studio/Tools/Rust/TokenizerResponse.cs b/app/MindWork AI Studio/Tools/Rust/TokenizerResponse.cs index f3ef0893..4eb89c6c 100644 --- a/app/MindWork AI Studio/Tools/Rust/TokenizerResponse.cs +++ b/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 = ""); diff --git a/app/MindWork AI Studio/Tools/Rust/TokenizerStatus.cs b/app/MindWork AI Studio/Tools/Rust/TokenizerStatus.cs deleted file mode 100644 index fd1b9dfc..00000000 --- a/app/MindWork AI Studio/Tools/Rust/TokenizerStatus.cs +++ /dev/null @@ -1,8 +0,0 @@ -namespace AIStudio.Tools.Rust; - -public enum TokenizerStatus -{ - UNAVAILABLE, - RUNNING, - AVAILABLE, -} diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs index dff7f369..61aa128c 100644 --- a/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs +++ b/app/MindWork AI Studio/Tools/Services/DataSourceEmbeddingService.Files.cs @@ -66,7 +66,7 @@ public sealed partial class DataSourceEmbeddingService { var segments = new List(); - 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 GetEmbeddingTokenCountAsync(EmbeddingProvider embeddingProvider, string text, CancellationToken token) { - var response = await rustService.GetTokenCount(embeddingProvider.Name, embeddingProvider.TokenizerPath, text, token); - if (response is { Success: true, Status: TokenizerStatus.AVAILABLE }) + var response = await rustService.GetTokenCount(embeddingProvider, text, token); + if (response is { Success: true }) return response.Value.TokenCount; var message = response?.Message ?? "No response was returned by the tokenizer service."; diff --git a/app/MindWork AI Studio/Tools/Services/DataSourceLocalRetrievalService.cs b/app/MindWork AI Studio/Tools/Services/DataSourceLocalRetrievalService.cs index e9dbfab7..04ab91ee 100644 --- a/app/MindWork AI Studio/Tools/Services/DataSourceLocalRetrievalService.cs +++ b/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( - embeddingProvider.Name, - embeddingProvider.TokenizerPath, - query, - token); - if (tokenCountResponse is not { Success: true, Status: TokenizerStatus.AVAILABLE }) + var tokenCountResponse = await rustService.GetTokenCount(embeddingProvider, query, token); + if (tokenCountResponse is not { Success: true }) { logger.LogWarning( "Skipping vector retrieval for data source '{DataSourceName}' ({DataSourceId}) because the token count for embedding provider '{EmbeddingProviderName}' could not be determined. Reason='{Reason}'.", diff --git a/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs b/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs index 38b1e01f..3598db6e 100644 --- a/app/MindWork AI Studio/Tools/Services/RustService.Retrieval.cs +++ b/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 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 StreamArbitraryFileDataWithTokenCounts( string path, - string providerName, - string tokenizerPath, + EmbeddingProvider embeddingProvider, [EnumeratorCancellation] CancellationToken token = default) { - await this.tokenizerLock.WaitAsync(token); - try + await foreach (var segment in this.StreamArbitraryFileDataCore(path, false, true, embeddingProvider.TokenizerPath, token)) { - var tokenizerResponse = await this.EnsureTokenizerCoreAsync(providerName, tokenizerPath); - if (tokenizerResponse is not { Success: true, Status: TokenizerStatus.AVAILABLE }) - { - var message = tokenizerResponse?.Message ?? "No response was returned by the tokenizer service."; - throw new InvalidOperationException($"Could not initialize tokenizer for provider '{providerName}'. {message}"); - } + if (segment.TokenCount is null) + throw new InvalidOperationException($"Rust did not return a token count for an extracted segment from '{path}' using provider '{embeddingProvider.Name}'."); - await foreach (var segment in this.StreamArbitraryFileDataCore(path, false, true, token)) - { - 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(); + yield return new(segment.Content, segment.TokenCount.Value); } } @@ -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); diff --git a/app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs b/app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs index 66da3766..7c62cd0f 100644 --- a/app/MindWork AI Studio/Tools/Services/RustService.Tokenizer.cs +++ b/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 GetTokenizerInfo(CancellationToken cancellationToken = default) - { - try - { - return await this.http.GetFromJsonAsync("/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 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(this.jsonRustSerializerOptions); } - public Task GetTokenCount(string text) - { - return this.GetTokenCountCoreAsync(text); - } + public Task GetTokenCount(AIStudio.Settings.Provider provider, string text, CancellationToken cancellationToken = default) => + this.GetTokenCount(provider.InstanceName, provider.TokenizerPath, text, cancellationToken); - public async Task GetTokenCount(string providerName, string path, string text, CancellationToken cancellationToken = default) - { - await this.tokenizerLock.WaitAsync(cancellationToken); - try - { - var tokenizerResponse = await this.EnsureTokenizerCoreAsync(providerName, path); - if (tokenizerResponse is not { Success: true, Status: TokenizerStatus.AVAILABLE }) - return tokenizerResponse; + public Task GetTokenCount(EmbeddingProvider provider, string text, CancellationToken cancellationToken = default) => + this.GetTokenCount(provider.Name, provider.TokenizerPath, text, cancellationToken); - return await this.GetTokenCountCoreAsync(text, cancellationToken); - } - finally - { - this.tokenizerLock.Release(); - } - } - - private async Task GetTokenCountCoreAsync(string text, CancellationToken cancellationToken = default) + private async Task GetTokenCount(string providerName, string tokenizerPath, string text, CancellationToken cancellationToken) { 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.hasInitializedTokenizer = false; + this.logger!.LogError("Failed to get the token count for provider '{ProviderName}': {StatusCode}", providerName, result.StatusCode); return CreateUnavailableTokenizerResponse("Error while getting token count from Rust service: "+result.StatusCode); } - var response = await result.Content.ReadFromJsonAsync(this.jsonRustSerializerOptions); - if (response is not { Status: TokenizerStatus.AVAILABLE }) - this.hasInitializedTokenizer = false; - - return response; - } - - public async Task 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(this.jsonRustSerializerOptions); - if (response is not { Success: true, Status: TokenizerStatus.AVAILABLE }) - this.hasInitializedTokenizer = false; - - return response; - } - - public async Task EnsureTokenizer(string providerName, string path) - { - await this.tokenizerLock.WaitAsync(); - try - { - return await this.EnsureTokenizerCoreAsync(providerName, path); - } - finally - { - this.tokenizerLock.Release(); - } - } - - private async Task 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; + return await result.Content.ReadFromJsonAsync(this.jsonRustSerializerOptions, cancellationToken); } } diff --git a/app/MindWork AI Studio/Tools/Services/RustService.cs b/app/MindWork AI Studio/Tools/Services/RustService.cs index 6f356444..6e979bb1 100644 --- a/app/MindWork AI Studio/Tools/Services/RustService.cs +++ b/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(); } diff --git a/runtime/src/file_data.rs b/runtime/src/file_data.rs index d783540d..4d116a58 100644 --- a/runtime/src/file_data.rs +++ b/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> { - self.token_count = Some(crate::tokenizer::get_segment_token_count(&self.content)?); + pub fn set_token_count(&mut self, tokenizer: &Tokenizer) -> std::result::Result<(), String> { + 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 @@ -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 Err(e) = chunk.set_token_count() { + if let Some(tokenizer) = tokenizer.as_deref() { + 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; } diff --git a/runtime/src/runtime_api.rs b/runtime/src/runtime_api.rs index a20128ff..d5c6e4d8 100644 --- a/runtime/src/runtime_api.rs +++ b/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)) diff --git a/runtime/src/tokenizer.rs b/runtime/src/tokenizer.rs index 6d554760..cd3ad19b 100644 --- a/runtime/src/tokenizer.rs +++ b/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>> = OnceLock::new(); +static TOKENIZERS: OnceLock>>> = OnceLock::new(); static DEFAULT_TOKENIZER_PATH: OnceLock = OnceLock::new(); -static TOKENIZER_STATUS: Lazy> = Lazy::new(|| Mutex::new(TokenizerStatusInfo::default())); -static TOKENIZER_OPERATION_LOCK: Lazy> = 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, -} +static TOKENIZER_STORAGE_LOCK: Mutex<()> = Mutex::new(()); #[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 { - 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) -> Json { - 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) } } -pub async fn set_tokenizer(_token: APIToken, payload: Json) -> Json { - 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 { - let _operation_guard = begin_tokenizer_operation()?; - 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 + validate_tokenizer_file(path) } -pub fn get_token_count(text: &str) -> Result { - get_token_count_internal(text, true) +pub fn get_token_count(path: &str, text: &str) -> Result { + let tokenizer = get_tokenizer(path)?; + get_token_count_internal(&tokenizer, text, true) } -pub fn get_segment_token_count(text: &str) -> Result { +pub fn get_segment_token_count(tokenizer: &Tokenizer, text: &str) -> Result { // 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 { +fn get_token_count_internal(tokenizer: &Tokenizer, text: &str, add_special_tokens: bool) -> Result { if text.trim().is_empty() { return Ok(0); } - let _operation_guard = begin_tokenizer_operation()?; - { - let status = TOKENIZER_STATUS.lock().unwrap(); - if status.status != TokenizerStatus::Available { - 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) + tokenizer + .encode(text, add_special_tokens) + .map(|encoding| encoding.len()) + .map_err(|e| format!("Failed to tokenize text: {e}")) } fn validate_tokenizer_file(path: &PathBuf) -> Result { @@ -293,7 +194,11 @@ fn handle_tokenizer_store(payload: &TokenizerStorage) -> Result 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> { - TOKENIZER.get_or_init(|| RwLock::new(None)) +fn tokenizer_cache() -> &'static RwLock>> { + TOKENIZERS.get_or_init(|| RwLock::new(HashMap::new())) } -fn begin_tokenizer_operation() -> Result, String> { - TOKENIZER_OPERATION_LOCK - .lock() - .map_err(|_| unavailable_with_status_update("Tokenizer operation lock is poisoned.")) -} +pub fn get_tokenizer(path: &str) -> Result, String> { + let resolved_path = resolve_tokenizer_path(path)?; + let tokenizer_path = fs::canonicalize(&resolved_path) + .map_err(|e| format!("Could not resolve tokenizer file '{}': {e}", resolved_path.display()))?; -fn set_tokenizer_available() { - let mut status = TOKENIZER_STATUS.lock().unwrap(); - status.status = TokenizerStatus::Available; - status.unavailable_reason = None; -} - -fn set_tokenizer_running() { - 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()), + if let Some(tokenizer) = tokenizer_cache() + .read() + .map_err(|_| "Tokenizer cache lock is poisoned.".to_string())? + .get(&tokenizer_path) + .cloned() + { + return Ok(tokenizer); } - set_tokenizer_unavailable(reason.clone()); - reason + let _storage_guard = TOKENIZER_STORAGE_LOCK + .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 {