tokengrams 0.3.0

Compute n-gram statistics and model language over pre-tokenized text corpora used to train large language models.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95

use funty::Unsigned;
use memmap2::{Mmap, MmapAsRawDesc, MmapMut};
use std::marker::PhantomData;
use std::ops::{Deref, DerefMut};

/// An immutable memory-mapped slice of unsigned integers
pub struct MmapSlice<T: Unsigned> {
    mmap: Mmap,
    _element_type: PhantomData<T>,
}

impl<T: Unsigned> MmapSlice<T> {
    pub fn new<F: MmapAsRawDesc>(file: F) -> std::io::Result<Self> {
        let raw = unsafe { Mmap::map(file)? };

        // Sanity check that the file size is a multiple of the element size.
        if raw.len() % std::mem::size_of::<T>() != 0 {
            Err(std::io::Error::new(
                std::io::ErrorKind::InvalidData,
                "File size is not a multiple of element size",
            ))
        } else {
            Ok(MmapSlice {
                mmap: raw,
                _element_type: PhantomData,
            })
        }
    }

    pub fn len(&self) -> usize {
        self.mmap.len() / std::mem::size_of::<T>()
    }

    pub fn as_slice<'a>(&'a self) -> &'a [T] {
        unsafe { std::slice::from_raw_parts(self.mmap.as_ptr() as *const T, self.len()) }
    }
}

impl<T: Unsigned> Deref for MmapSlice<T> {
    type Target = [T];

    fn deref(&self) -> &[T] {
        self.as_slice()
    }
}

/// A mutable memory-mapped slice of unsigned integers
pub struct MmapSliceMut<T: Unsigned> {
    mmap: MmapMut,
    _element_type: PhantomData<T>,
}

impl<T: Unsigned> MmapSliceMut<T> {
    pub fn new<F: MmapAsRawDesc>(file: F) -> std::io::Result<Self> {
        let raw = unsafe { MmapMut::map_mut(file)? };

        Ok(MmapSliceMut {
            mmap: raw,
            _element_type: PhantomData,
        })
    }

    pub fn len(&self) -> usize {
        self.mmap.len() / std::mem::size_of::<T>()
    }

    pub fn as_slice<'a>(&'a self) -> &'a [T] {
        unsafe { std::slice::from_raw_parts(self.mmap.as_ptr() as *const T, self.len()) }
    }

    pub fn as_slice_mut<'a>(&'a mut self) -> &'a mut [T] {
        unsafe { std::slice::from_raw_parts_mut(self.mmap.as_mut_ptr() as *mut T, self.len()) }
    }

    pub fn into_read_only(self) -> std::io::Result<MmapSlice<T>> {
        Ok(MmapSlice {
            mmap: self.mmap.make_read_only()?,
            _element_type: PhantomData,
        })
    }
}

impl<T: Unsigned> Deref for MmapSliceMut<T> {
    type Target = [T];

    fn deref(&self) -> &[T] {
        self.as_slice()
    }
}

impl<T: Unsigned> DerefMut for MmapSliceMut<T> {
    fn deref_mut(&mut self) -> &mut [T] {
        self.as_slice_mut()
    }
}