File size: 708 Bytes
b7a913f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
from model_registry import validate_model_id
import model_service

def count_tokens(model_id: str, text: str) -> int:
    """
    Get actual token count using the model's tokenizer.
    If the model is not loaded, it loads the tokenizer-only temporarily.
    """
    validate_model_id(model_id)
    
    # Check if the active tokenizer is the correct one
    if model_service.active_model_id == model_id and model_service.active_tokenizer is not None:
        tokenizer = model_service.active_tokenizer
    else:
        # Load tokenizer dynamically
        tokenizer = model_service.load_tokenizer_only(model_id)
        
    tokens = tokenizer.encode(text, add_special_tokens=False)
    return len(tokens)