lance_tokenizer/
tokenizer_api.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3// SPDX-License-Identifier: MIT
4// Adapted from Tantivy v0.24.2 tokenizer API.
5// Copyright (c) 2017-present Tantivy contributors.
6
7use std::borrow::{Borrow, BorrowMut};
8use std::ops::{Deref, DerefMut};
9
10/// Token emitted by a tokenizer.
11#[derive(Debug, Clone, Eq, PartialEq)]
12pub struct Token {
13    /// Byte offset of the first character.
14    pub offset_from: usize,
15    /// Byte offset after the last character.
16    pub offset_to: usize,
17    /// Logical token position.
18    pub position: usize,
19    /// Token text.
20    pub text: String,
21    /// Position length measured in original tokens.
22    pub position_length: usize,
23}
24
25impl Default for Token {
26    fn default() -> Self {
27        Self {
28            offset_from: 0,
29            offset_to: 0,
30            position: usize::MAX,
31            text: String::new(),
32            position_length: 1,
33        }
34    }
35}
36
37impl Token {
38    /// Reset the token to its default state.
39    pub fn reset(&mut self) {
40        self.offset_from = 0;
41        self.offset_to = 0;
42        self.position = usize::MAX;
43        self.text.clear();
44        self.position_length = 1;
45    }
46}
47
48/// Tokenizer splits text into a token stream.
49pub trait Tokenizer: 'static + Clone + Send + Sync {
50    /// Stream type emitted by the tokenizer.
51    type TokenStream<'a>: TokenStream;
52
53    /// Create a token stream for the provided text.
54    fn token_stream<'a>(&'a mut self, text: &'a str) -> Self::TokenStream<'a>;
55}
56
57/// Token stream object-safe wrapper.
58pub struct BoxTokenStream<'a> {
59    token_stream: Box<dyn TokenStream + 'a>,
60    error: Option<String>,
61}
62
63impl<'a> BoxTokenStream<'a> {
64    /// Wrap a token stream.
65    pub fn new<T: TokenStream + 'a>(token_stream: T) -> Self {
66        Self {
67            token_stream: Box::new(token_stream),
68            error: None,
69        }
70    }
71
72    /// Attach an error that is returned before the first token.
73    pub fn with_error(mut self, error: impl Into<String>) -> Self {
74        self.error = Some(error.into());
75        self
76    }
77
78    /// Return the next token or the deferred stream error.
79    pub fn try_next(&mut self) -> Result<Option<&Token>, String> {
80        if let Some(error) = self.error.take() {
81            return Err(error);
82        }
83        Ok(self.token_stream.next())
84    }
85}
86
87impl TokenStream for BoxTokenStream<'_> {
88    fn advance(&mut self) -> bool {
89        self.token_stream.advance()
90    }
91
92    fn token(&self) -> &Token {
93        self.token_stream.token()
94    }
95
96    fn token_mut(&mut self) -> &mut Token {
97        self.token_stream.token_mut()
98    }
99}
100
101impl<'a> Deref for BoxTokenStream<'a> {
102    type Target = dyn TokenStream + 'a;
103
104    fn deref(&self) -> &Self::Target {
105        &*self.token_stream
106    }
107}
108
109impl DerefMut for BoxTokenStream<'_> {
110    fn deref_mut(&mut self) -> &mut Self::Target {
111        &mut *self.token_stream
112    }
113}
114
115impl<'a> TokenStream for Box<dyn TokenStream + 'a> {
116    fn advance(&mut self) -> bool {
117        let token_stream: &mut dyn TokenStream = self.borrow_mut();
118        token_stream.advance()
119    }
120
121    fn token(&self) -> &Token {
122        let token_stream: &(dyn TokenStream + 'a) = self.borrow();
123        token_stream.token()
124    }
125
126    fn token_mut(&mut self) -> &mut Token {
127        let token_stream: &mut (dyn TokenStream + 'a) = self.borrow_mut();
128        token_stream.token_mut()
129    }
130}
131
132/// Consumable token stream.
133pub trait TokenStream {
134    /// Advance to the next token.
135    fn advance(&mut self) -> bool;
136
137    /// Access the current token.
138    fn token(&self) -> &Token;
139
140    /// Mutate the current token.
141    fn token_mut(&mut self) -> &mut Token;
142
143    /// Iterate to the next token and return it.
144    fn next(&mut self) -> Option<&Token> {
145        if self.advance() {
146            Some(self.token())
147        } else {
148            None
149        }
150    }
151
152    /// Consume the remaining stream into the provided sink.
153    fn process(&mut self, sink: &mut dyn FnMut(&Token)) {
154        while self.advance() {
155            sink(self.token());
156        }
157    }
158}
159
160/// Filter that wraps a tokenizer with additional token-processing behavior.
161pub trait TokenFilter: 'static + Send + Sync {
162    /// Tokenizer produced by this filter.
163    type Tokenizer<T: Tokenizer>: Tokenizer;
164
165    /// Wrap the tokenizer.
166    fn transform<T: Tokenizer>(self, tokenizer: T) -> Self::Tokenizer<T>;
167}
lance_tokenizer/tokenizer_api.rs

lance_tokenizer/
tokenizer_api.rs
