Skip to main content

token_dict/
dict.rs

1impl AsMut<Self> for TokenDict{
2	fn as_mut(&mut self)->&mut Self{self}
3}
4impl AsRef<Self> for TokenDict{
5	fn as_ref(&self)->&Self{self}
6}
7impl Default for TokenDict{
8	fn default()->Self{
9		let x:[Vec<u8>;0]=[];
10		x.into_iter().collect()
11	}
12}
13impl DoubleEndedIterator for DictIntoIter{
14	fn next_back(&mut self)->Option<Self::Item>{
15		self.range.next_back().map(|n|if n<256{SINGLE_TOKENS[n].clone()}else{self.tokens[n-256].clone()})
16	}
17	fn nth_back(&mut self,n:usize)->Option<Self::Item>{
18		self.range.nth_back(n).map(|n|if n<256{SINGLE_TOKENS[n].clone()}else{self.tokens[n-256].clone()})
19	}
20	fn rfold<B,F:FnMut(B,Self::Item)->B>(self,mut init:B,mut f:F)->B{
21		let (range,tokens)=(self.range,self.tokens);
22		let (start,stop)=(range.start,range.end);
23
24		init=tokens[start.saturating_sub(256)..stop.saturating_sub(256)].iter().cloned().rfold(init,&mut f);
25		SINGLE_TOKENS[start.min(256)..stop.min(256)].iter().cloned().rfold(init,f)
26	}
27}
28impl ExactSizeIterator for DictIntoIter{
29	fn len(&self)->usize{self.range.len()}
30}
31impl Index<u32> for TokenDict{
32	fn index(&self,ix:u32)->&Self::Output{
33		let ix=ix as usize;
34		if ix<256{&SINGLE_BYTES[ix..ix+1]}else{&self.tokens[ix-256]}
35	}
36	type Output=[u8];
37}
38impl Index<usize> for TokenDict{
39	fn index(&self,ix:usize)->&Self::Output{
40		if ix<256{&SINGLE_TOKENS[ix]}else{&self.tokens[ix-256]}
41	}
42	type Output=Token;
43}
44impl IntoIterator for TokenDict{
45	fn into_iter(self)->Self::IntoIter{
46		DictIntoIter{range:0..self.len(),tokens:self.tokens}
47	}
48	type IntoIter=DictIntoIter;
49	type Item=Token;
50}
51impl Iterator for DictIntoIter{
52	fn count(self)->usize{self.range.count()}
53	fn fold<B,F:FnMut(B,Self::Item)->B>(self,mut init:B,mut f:F)->B{
54		let (range,tokens)=(self.range,self.tokens);
55		let (start,stop)=(range.start,range.end);
56
57		init=SINGLE_TOKENS[start.min(256)..stop.min(256)].iter().cloned().fold(init,&mut f);
58		tokens[start.saturating_sub(256)..stop.saturating_sub(256)].iter().cloned().fold(init,f)
59	}
60	fn last(mut self)->Option<Self::Item>{self.next_back()}
61	fn next(&mut self)->Option<Self::Item>{
62		self.range.next().map(|n|if n<256{SINGLE_TOKENS[n].clone()}else{self.tokens[n-256].clone()})
63	}
64	fn nth(&mut self,n:usize)->Option<Self::Item>{
65		self.range.nth(n).map(|n|if n<256{SINGLE_TOKENS[n].clone()}else{self.tokens[n-256].clone()})
66	}
67	fn size_hint(&self)->(usize,Option<usize>){self.range.size_hint()}
68	type Item=Token;
69}
70#[cfg(feature="serial")]
71impl Serialize for TokenDict{
72	fn serialize<S:Serializer>(&self,serializer:S)->Result<S::Ok,S::Error>{
73		let data:Vec<Vec<u8>>=self.iter().map(|t|t.to_vec()).collect();
74		data.serialize(serializer)
75	}
76}
77impl TokenDict{
78	/// decodes the tokens into bytes
79	pub fn detokenize<I:IntoIterator>(&self,tokens:I)->Detokenization<I::IntoIter> where I::Item:Val<u32>{
80		Detokenization{inner:tokens.into_iter().fuse(),maxtokenlen:self.maxtokenlen,position:1,tokenid:0,tokens:self.tokens.clone()}
81	}
82	/// creates an iterator over tokens
83	pub fn detoken_iter<I:IntoIterator>(&self,tokens:I)->impl Iterator<Item=Token> where I::Item:Val<u32>{
84		let tokenizer=self.clone();
85		tokens.into_iter().map(move|id|tokenizer[id.val() as usize].clone())
86	}
87	/// decodes the tokens into chars, replacing invalid unicode with replacement character
88	pub fn detokenize_str<I:IntoIterator>(&self,tokens:I)->impl Iterator<Item=char> where I::Item:Val<u32>{
89		UTF8CharIter::from(self.detokenize(tokens)).map(|r|if let Ok(c)=r{c}else{char::REPLACEMENT_CHARACTER})
90	}
91	/// decodes the tokens into chars, replacing invalid unicode with replacement character
92	pub fn detokenize_string<I:IntoIterator>(&self,tokens:I)->String where I::Item:Val<u32>{self.detokenize_str(tokens).collect()}
93	/// accumulates frequencies of each token in the text as if they were tokenized by this tokenizer
94	pub fn frequencies<I:IntoIterator,O:Into<Option<Vec<usize>>>>(&self,data:I,freq:O)->Vec<usize> where I::Item:Val<u8>{
95		let mut freq=freq.into().unwrap_or_default();
96		if freq.len()<self.len(){freq.resize(self.len(),0)}
97
98		for t in self.tokenize(data){freq[t as usize]+=1}
99		freq
100	}
101	/// gets an id for the token if it is in the dictionary
102	pub fn get_id(&self,token:&[u8])->Option<u32>{self.ids[*token.get(0)? as usize].get(token.iter().copied().skip(1)).copied()}
103	/// decodes the tokens into bytes
104	pub fn into_detokenize<I:IntoIterator>(self,tokens:I)->Detokenization<I::IntoIter> where I::Item:Val<u32>{
105		Detokenization{inner:tokens.into_iter().fuse(),maxtokenlen:self.maxtokenlen,position:1,tokenid:0,tokens:self.tokens}
106	}
107	/// creates an iterator over tokens
108	pub fn into_detoken_iter<I:IntoIterator>(self,tokens:I)->impl Iterator<Item=Token> where I::Item:Val<u32>{
109		let tokenizer=self.clone();
110		tokens.into_iter().map(move|id|tokenizer[id.val() as usize].clone())
111	}
112	/// creates an iterator over tokens
113	pub fn into_token_iter<I:IntoIterator>(self,bytes:I)->impl Iterator<Item=Token> where I::Item:Val<u8>{
114		let tokenizer=self.clone();
115		self.tokenize(bytes).map(move|id|tokenizer[id as usize].clone())
116	}
117	/// converts the bytes to tokens
118	pub fn into_tokenize<I:IntoIterator>(self,bytes:I)->Tokenization<I::IntoIter> where I::Item:Val<u8>{
119		Tokenization{ids:self.ids,inner:bytes.into_iter().fuse(),state:VecDeque::with_capacity(self.maxtokenlen)}
120	}
121	/// returns an interator over the possible tokens generated by this tokenizer
122	pub fn iter(&self)->DictIter<'_>{
123		DictIter{range:0..self.len(),tokens:&self.tokens}
124	}
125	/// returns the number of possible token ids generated by this tokenizer
126	pub fn len(&self)->usize{self.tokens.len()+256}
127	/// finds token pairs and returns new tokens of them mapped to their frequencies with ids as if they were added to this dictionary. Tokens with ids within the current dictionary will have those ids
128	pub fn pairs<I:IntoIterator,O:Into<Option<HashMap<Token,usize>>>>(&self,data:I,freq:O)->HashMap<Token,usize> where I::Item:Val<u8>{
129		let mut freq=freq.into().unwrap_or_default();
130		let mut nexttokenid=freq.keys().map(|t|t.id()+1).chain([self.len() as u32]).max().unwrap();
131		let mut previous:Option<u32>=None;
132		let mut temp:Vec<u8>=Vec::new();
133
134		self.tokenize(data).for_each(|id|{
135			if let Some(previous)=previous{
136				temp.clear();
137				temp.extend(self[previous as usize].clone());
138				temp.extend(self[id as usize].clone());
139
140				if let Some(f)=freq.get_mut(temp.as_slice()){
141					*f+=1
142				}else{
143					let newid=if let Some(id)=self.get_id(&temp){id}else{post_inc!(nexttokenid)};
144					let token=Token::new(newid,Some(Arc::from(temp.as_slice())));
145
146					freq.insert(token,1);
147				}
148			}
149			previous=Some(id);
150		});
151		freq
152	}
153	/// adds the token to the dictionary
154	pub fn push<A:AsRef<[u8]>>(&mut self,token:A){self.extend(Some(token))}
155	/// converts the string to a token vec
156	pub fn string_to_tokens<S:?Sized+AsRef<str>>(&self,input:&S)->Vec<u32>{self.tokenize(input.as_ref().as_bytes()).collect()}
157	/// creates an iterator over tokens
158	pub fn token_iter<I:IntoIterator>(&self,bytes:I)->impl Iterator<Item=Token> where I::Item:Val<u8>{
159		let tokenizer=self.clone();
160		self.tokenize(bytes).map(move|id|tokenizer[id as usize].clone())
161	}
162	/// converts the bytes to tokens
163	pub fn tokenize<I:IntoIterator>(&self,bytes:I)->Tokenization<I::IntoIter> where I::Item:Val<u8>{
164		Tokenization{ids:self.ids.clone(),inner:bytes.into_iter().fuse(),state:VecDeque::with_capacity(self.maxtokenlen)}
165	}
166	/// converts the string to tokens
167	pub fn tokenize_str<'a,S:?Sized+AsRef<str>>(&self,input:&'a S)->Tokenization<SliceIter<'a,u8>>{self.tokenize(input.as_ref().as_bytes())}
168	/// converts the string to tokens
169	pub fn tokenize_string(&self,input:String)->Tokenization<VecIntoIter<u8>>{self.tokenize(Vec::from(input))}
170	/// converts the token vec to string
171	pub fn tokens_to_string<V:?Sized+AsRef<[u32]>>(&self,input:&V)->String{String::from_utf8_lossy(&self.detokenize(input.as_ref()).collect::<Vec<u8>>()).to_string()}
172}
173#[cfg(feature="serial")]
174impl<'a> Deserialize<'a> for TokenDict{
175	fn deserialize<D:Deserializer<'a>>(deserializer:D)->Result<Self,D::Error>{
176		let data:Vec<Vec<u8>>=Deserialize::deserialize(deserializer)?;
177		Ok(data.into_iter().collect())
178	}
179}
180impl<'a> DoubleEndedIterator for DictIter<'a>{
181	fn next_back(&mut self)->Option<Self::Item>{
182		self.range.next_back().map(|n|if n<256{&SINGLE_TOKENS[n]}else{&self.tokens[n-256]})
183	}
184	fn nth_back(&mut self,n:usize)->Option<Self::Item>{
185		self.range.nth_back(n).map(|n|if n<256{&SINGLE_TOKENS[n]}else{&self.tokens[n-256]})
186	}
187	fn rfold<B,F:FnMut(B,Self::Item)->B>(self,mut init:B,mut f:F)->B{
188		let (range,tokens)=(self.range,self.tokens);
189		let (start,stop)=(range.start,range.end);
190
191		init=tokens[start.saturating_sub(256)..stop.saturating_sub(256)].iter().rfold(init,&mut f);
192		SINGLE_TOKENS[start.min(256)..stop.min(256)].iter().rfold(init,f)
193	}
194}
195impl<'a> ExactSizeIterator for DictIter<'a>{
196	fn len(&self)->usize{self.range.len()}
197}
198impl<'a> Iterator for DictIter<'a>{
199	fn count(self)->usize{self.range.count()}
200	fn fold<B,F:FnMut(B,Self::Item)->B>(self,mut init:B,mut f:F)->B{
201		let (range,tokens)=(self.range,self.tokens);
202		let (start,stop)=(range.start,range.end);
203
204		init=SINGLE_TOKENS[start.min(256)..stop.min(256)].iter().fold(init,&mut f);
205		tokens[start.saturating_sub(256)..stop.saturating_sub(256)].iter().fold(init,f)
206	}
207	fn last(mut self)->Option<Self::Item>{self.next_back()}
208	fn next(&mut self)->Option<Self::Item>{
209		self.range.next().map(|n|if n<256{&SINGLE_TOKENS[n]}else{&self.tokens[n-256]})
210	}
211	fn nth(&mut self,n:usize)->Option<Self::Item>{
212		self.range.nth(n).map(|n|if n<256{&SINGLE_TOKENS[n]}else{&self.tokens[n-256]})
213	}
214	fn size_hint(&self)->(usize,Option<usize>){self.range.size_hint()}
215	type Item=&'a Token;
216}
217impl<A:AsRef<[u8]>> Extend<A> for TokenDict{
218	fn extend<I:IntoIterator<Item=A>>(&mut self,iter:I){
219		let (ids,tokens)=(Arc::make_mut(&mut self.ids),Arc::make_mut(&mut self.tokens));
220		let maxtokenlen=&mut self.maxtokenlen;
221
222		iter.into_iter().filter(|a|a.as_ref().len()>1).for_each(|a|{
223			let id=u32::try_from(tokens.len()+256).unwrap();
224			let token:Arc<[u8]>=Arc::from(a.as_ref());
225
226			ids[token[0] as usize].insert(token.iter().copied().skip(1),id);
227			*maxtokenlen=(*maxtokenlen).max(token.len());
228			tokens.push(Token::new(id,Some(token)))
229		});
230	}
231}
232impl<A:AsRef<[u8]>> FromIterator<A> for TokenDict{
233	fn from_iter<I:IntoIterator<Item=A>>(iter:I)->Self{
234		let mut maxtokenlen=1;
235		let mut ids:[Trie<_,_>;256]=std::array::from_fn(|_|Trie::new());
236		let tokens:Vec<Token>=iter.into_iter().filter(|t|t.as_ref().len()>1).enumerate().map(|(n,t)|{
237			let id=u32::try_from(n+256).unwrap();
238			let token:Arc<[u8]>=Arc::from(t.as_ref());
239
240			ids[token[0] as usize].insert(token.iter().copied().skip(1),id);
241			maxtokenlen=maxtokenlen.max(token.len());
242
243			Token::new(id,Some(token))
244		}).collect();
245
246		let (ids,tokens)=(Arc::new(ids),Arc::new(tokens));
247		Self{ids,maxtokenlen,tokens}
248	}
249}
250impl<I:Iterator> Iterator for Detokenization<I> where I::Item:Val<u32>{
251	fn fold<B,F:FnMut(B,Self::Item)->B>(self,init:B,mut f:F)->B{
252		self.inner.map(Val::val).fold(init,|acc,tokenid|if tokenid<256{f(acc,tokenid as u8)}else{self.tokens[tokenid as usize].iter().fold(acc,|acc,&b|f(acc,b))})
253	}
254	fn next(&mut self)->Option<u8>{
255		let (inner,position)=(&mut self.inner,&mut self.position);
256		let tokenid=&mut self.tokenid;
257		let tokens=&self.tokens;
258
259		if let Some(b)=if *tokenid<256{(*position==0).then_some(*tokenid as u8)}else{tokens[*tokenid as usize-256].get(*position).map(|&b|b)}{
260			*position+=1;
261			b
262		}else{
263			*position=1;
264			*tokenid=inner.map(Val::val).next()?;
265			if *tokenid<256{*tokenid as u8}else{tokens[*tokenid as usize-256][0]}
266		}.into()
267	}
268	fn size_hint(&self)->(usize,Option<usize>){
269		let (lowertokens,uppertokens)=self.inner.size_hint();
270		let maxtoken=self.maxtokenlen;
271
272		(lowertokens,uppertokens.map(|h|h*maxtoken))
273	}
274	type Item=u8;
275}
276impl<I:Iterator> Iterator for Tokenization<I> where I::Item:Val<u8>{
277	fn next(&mut self)->Option<u32>{
278		let (inner,state)=(&mut self.inner,&mut self.state);
279		let ids=&self.ids;
280		state.extend(inner.map(Val::val).take(state.capacity()-state.len()));
281		if state.len()==0{return None}
282
283		let (tokenlen,&tokenid)=if let Some(t)=ids[state[0] as usize].find_longest_prefix_len(state.iter().copied().skip(1)).filter(|(tokenlen,_tokenid)|*tokenlen>0){t}else{return Some(state.pop_front().unwrap() as u32)};
284		state.drain(..tokenlen+1);
285		return Some(tokenid)
286	}
287	fn size_hint(&self)->(usize,Option<usize>){
288		let (lowerbytes,upperbytes)=self.inner.size_hint();
289		let maxtoken=self.state.capacity();
290		let statelen=self.state.len();
291
292		((lowerbytes+statelen).div_ceil(maxtoken),upperbytes.map(|b|b+statelen))
293	}
294	type Item=u32;
295}
296
297macro_rules! post_inc {
298	($e:expr) => {{
299		let old = $e;
300		$e += 1;
301		old
302	}};
303}
304
305#[cfg(test)]
306mod tests{
307	#[test]
308	fn tokenizer_iter(){
309		let tokenizer:TokenDict=["aa","bb","cc"].into_iter().collect();
310		let t2:TokenDict=tokenizer.iter().collect();
311		assert_eq!(tokenizer.tokenize_str("ccaabb").collect::<Vec<_>>(),t2.tokenize_str("ccaabb").collect::<Vec<_>>());
312	}
313	#[test]
314	fn bytes_only(){
315		let teststring="oishsoghohhduihahdufghud";
316		let tokenizer=TokenDict::default();
317		let tokens:Vec<u32>=tokenizer.tokenize_str(teststring).collect();
318		let detokens:Vec<u8>=tokenizer.detokenize(tokens).collect();
319
320		assert_eq!(detokens.as_slice(),teststring.as_bytes());
321	}
322	#[test]
323	fn there_are_tokens_yay(){
324		let teststring="there are tokens! yay";
325		let tokenizer:TokenDict=["there","are","tokens","yay"].into_iter().collect();
326		let tokens:Vec<u32>=tokenizer.tokenize(teststring.bytes()).collect();
327		let detokens:Vec<u8>=tokenizer.detokenize(&tokens).collect();
328
329		assert_eq!(tokens.len(),8);
330		assert_eq!(detokens.as_slice(),teststring.as_bytes());
331	}
332	#[test]
333	fn test_default_token_dict_detokenize_empty() {
334		let dict = TokenDict::default();
335		let inp:Vec<u32>=vec![];
336		let out: Vec<u8> = dict.detokenize(inp).collect();
337		assert!(out.is_empty(), "Detokenizing an empty input should yield no bytes");
338	}
339	use super::*;
340}
341
342pub (crate) const SINGLE_TOKENS:&[Token;256]=&[Token::single(0),Token::single(1),Token::single(2),Token::single(3),Token::single(4),Token::single(5),Token::single(6),Token::single(7),Token::single(8),Token::single(9),Token::single(10),Token::single(11),Token::single(12),Token::single(13),Token::single(14),Token::single(15),Token::single(16),Token::single(17),Token::single(18),Token::single(19),Token::single(20),Token::single(21),Token::single(22),Token::single(23),Token::single(24),Token::single(25),Token::single(26),Token::single(27),Token::single(28),Token::single(29),Token::single(30),Token::single(31),Token::single(32),Token::single(33),Token::single(34),Token::single(35),Token::single(36),Token::single(37),Token::single(38),Token::single(39),Token::single(40),Token::single(41),Token::single(42),Token::single(43),Token::single(44),Token::single(45),Token::single(46),Token::single(47),Token::single(48),Token::single(49),Token::single(50),Token::single(51),Token::single(52),Token::single(53),Token::single(54),Token::single(55),Token::single(56),Token::single(57),Token::single(58),Token::single(59),Token::single(60),Token::single(61),Token::single(62),Token::single(63),Token::single(64),Token::single(65),Token::single(66),Token::single(67),Token::single(68),Token::single(69),Token::single(70),Token::single(71),Token::single(72),Token::single(73),Token::single(74),Token::single(75),Token::single(76),Token::single(77),Token::single(78),Token::single(79),Token::single(80),Token::single(81),Token::single(82),Token::single(83),Token::single(84),Token::single(85),Token::single(86),Token::single(87),Token::single(88),Token::single(89),Token::single(90),Token::single(91),Token::single(92),Token::single(93),Token::single(94),Token::single(95),Token::single(96),Token::single(97),Token::single(98),Token::single(99),Token::single(100),Token::single(101),Token::single(102),Token::single(103),Token::single(104),Token::single(105),Token::single(106),Token::single(107),Token::single(108),Token::single(109),Token::single(110),Token::single(111),Token::single(112),Token::single(113),Token::single(114),Token::single(115),Token::single(116),Token::single(117),Token::single(118),Token::single(119),Token::single(120),Token::single(121),Token::single(122),Token::single(123),Token::single(124),Token::single(125),Token::single(126),Token::single(127),Token::single(128),Token::single(129),Token::single(130),Token::single(131),Token::single(132),Token::single(133),Token::single(134),Token::single(135),Token::single(136),Token::single(137),Token::single(138),Token::single(139),Token::single(140),Token::single(141),Token::single(142),Token::single(143),Token::single(144),Token::single(145),Token::single(146),Token::single(147),Token::single(148),Token::single(149),Token::single(150),Token::single(151),Token::single(152),Token::single(153),Token::single(154),Token::single(155),Token::single(156),Token::single(157),Token::single(158),Token::single(159),Token::single(160),Token::single(161),Token::single(162),Token::single(163),Token::single(164),Token::single(165),Token::single(166),Token::single(167),Token::single(168),Token::single(169),Token::single(170),Token::single(171),Token::single(172),Token::single(173),Token::single(174),Token::single(175),Token::single(176),Token::single(177),Token::single(178),Token::single(179),Token::single(180),Token::single(181),Token::single(182),Token::single(183),Token::single(184),Token::single(185),Token::single(186),Token::single(187),Token::single(188),Token::single(189),Token::single(190),Token::single(191),Token::single(192),Token::single(193),Token::single(194),Token::single(195),Token::single(196),Token::single(197),Token::single(198),Token::single(199),Token::single(200),Token::single(201),Token::single(202),Token::single(203),Token::single(204),Token::single(205),Token::single(206),Token::single(207),Token::single(208),Token::single(209),Token::single(210),Token::single(211),Token::single(212),Token::single(213),Token::single(214),Token::single(215),Token::single(216),Token::single(217),Token::single(218),Token::single(219),Token::single(220),Token::single(221),Token::single(222),Token::single(223),Token::single(224),Token::single(225),Token::single(226),Token::single(227),Token::single(228),Token::single(229),Token::single(230),Token::single(231),Token::single(232),Token::single(233),Token::single(234),Token::single(235),Token::single(236),Token::single(237),Token::single(238),Token::single(239),Token::single(240),Token::single(241),Token::single(242),Token::single(243),Token::single(244),Token::single(245),Token::single(246),Token::single(247),Token::single(248),Token::single(249),Token::single(250),Token::single(251),Token::single(252),Token::single(253),Token::single(254),Token::single(255)];
343#[derive(Clone,Debug)]
344/// a simple dictionary based detokenizer
345pub struct Detokenization<I:Iterator> where I::Item:Val<u32>{inner:Fuse<I>,maxtokenlen:usize,position:usize,tokenid:u32,tokens:Arc<Vec<Token>>}
346#[derive(Clone,Debug)]
347/// iterator for TokenDict
348pub struct DictIter<'a>{range:Range<usize>,tokens:&'a [Token]}
349#[derive(Clone,Debug)]
350/// intoiterator for TokenDict
351pub struct DictIntoIter{range:Range<usize>,tokens:Arc<Vec<Token>>}
352#[derive(Clone,Debug)]
353/// a simple dictionary based tokenizer dictionary where single bytes implicitly form the lower 256 ids. It's reference counted and cheap to clone
354pub struct TokenDict{ids:Arc<[Trie<u8,u32>;256]>,maxtokenlen:usize,tokens:Arc<Vec<Token>>}
355#[derive(Clone,Debug)]
356/// a simple dictionary based tokenizer iterator
357pub struct Tokenization<I:Iterator> where I::Item:Val<u8>{ids:Arc<[Trie<u8,u32>;256]>,inner:Fuse<I>,state:VecDeque<u8>}
358
359use crate::{Token,UTF8CharIter,Val,token::SINGLE_BYTES};
360use post_inc;
361use ptrie::Trie;
362#[cfg(feature="serial")]
363use serde::{Deserialize,Deserializer,Serialize,Serializer};
364use std::{
365	collections::{HashMap,VecDeque},iter::{Extend,Fuse},ops::{Index,Range},slice::Iter as SliceIter,sync::Arc,vec::IntoIter as VecIntoIter,
366};