Coverage for src/semware/utils/tokenizer.py: 100%

1"""Text tokenization utilities for handling context length limits.""" 

2 

3 

4import tiktoken 

5from loguru import logger 

6 

7 

8class TextTokenizer: 

9 """Handles text tokenization and batching for embedding generation.""" 

10 

11 def __init__(self, encoding_name: str = "cl100k_base"): 

12 """Initialize tokenizer. 

13 

14 Args: 

15 encoding_name: Name of the tokenizer encoding to use 

16 """ 

17 self.encoding = tiktoken.get_encoding(encoding_name) 

18 logger.info(f"Initialized tokenizer with encoding: {encoding_name}") 

19 

20 def count_tokens(self, text: str) -> int: 

21 """Count the number of tokens in the text. 

22 

23 Args: 

24 text: Input text 

25 

26 Returns: 

27 Number of tokens 

28 """ 

29 return len(self.encoding.encode(text)) 

30 

31 def batch_text(self, text: str, max_tokens: int = 2000) -> list[str]: 

32 """Split text into batches based on token count. 

33 

34 Args: 

35 text: Input text to batch 

36 max_tokens: Maximum tokens per batch 

37 

38 Returns: 

39 List of text batches 

40 """ 

41 if not text.strip(): 

42 return [] 

43 

44 tokens = self.encoding.encode(text) 

45 total_tokens = len(tokens) 

46 

47 if total_tokens <= max_tokens: 

48 return [text] 

49 

50 # Split tokens into batches 

51 batches = [] 

52 for i in range(0, total_tokens, max_tokens): 

53 batch_tokens = tokens[i : i + max_tokens] 

54 batch_text = self.encoding.decode(batch_tokens) 

55 batches.append(batch_text) 

56 

57 logger.debug( 

58 f"Split text into {len(batches)} batches ({total_tokens} total tokens)" 

59 ) 

60 return batches 

61 

62 def batch_texts(self, texts: list[str], max_tokens: int = 2000) -> list[list[str]]: 

63 """Batch multiple texts. 

64 

65 Args: 

66 texts: List of texts to batch 

67 max_tokens: Maximum tokens per batch 

68 

69 Returns: 

70 List of batches for each text 

71 """ 

72 return [self.batch_text(text, max_tokens) for text in texts] 

73 

74 

75# Global tokenizer instance 

76tokenizer = TextTokenizer()