Address review feedback

ggml/src/lib.rs

@@ -406,7 +406,7 @@ impl Tensor {
         }
     }
 
-    fn with_alive_ctx<U>(&self, f: impl Fn() -> U) -> U {
+    fn with_alive_ctx<U>(&self, mut f: impl FnMut() -> U) -> U {
         if let Some(_ctx) = self.ctx.upgrade() {
             f()
         } else {

llama-rs/src/convert.rs

@@ -129,9 +129,6 @@ fn write_header(fout: &mut File, hparams: &Hyperparameters) -> Result<(), String
 fn write_tokens(file: &mut File, vocab: &Vocabulary) -> Result<(), String> {
     let mut values: Vec<u8> = vec![];
     for (i, token) in vocab.id_to_token.iter().enumerate() {
-        // TODO: Not sure what the behaviour should be if the token is not valid UTF-8.
-        //
-        // Switching to the HF tokenizer should fix this.
         let text = if let Ok(token) = std::str::from_utf8(token) {
             match token {
                 _ if token.contains("<unk>") => " \u{2047} ".as_bytes().to_vec(),

llama-rs/src/lib.rs

@@ -570,7 +570,7 @@ impl Model {
     pub fn load(
         path: impl AsRef<Path>,
         n_context_tokens: usize,
-        load_progress_callback: impl Fn(LoadProgress),
+        mut load_progress_callback: impl FnMut(LoadProgress),
     ) -> Result<(Model, Vocabulary), LoadError> {
         use std::fs::File;
         use std::io::BufReader;
@@ -1768,7 +1768,21 @@ impl TokenUtf8Buffer {
                 self.0 = vec![];
                 Some(out)
             }
-            Err(..) => None,
+            Err(..) => {
+                for i in 1..self.0.len() {
+                    let slice = &self.0[i..];
+                    if slice.is_empty() {
+                        break;
+                    }
+
+                    if let Ok(s) = std::str::from_utf8(slice) {
+                        let out = s.to_owned();
+                        self.0 = vec![];
+                        return Some(out);
+                    }
+                }
+                None
+            }
         }
     }
 
@@ -1783,3 +1797,30 @@ impl TokenUtf8Buffer {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_valid_utf8() {
+        let mut buffer = TokenUtf8Buffer::new();
+        assert_eq!(buffer.push(b"hello").as_deref(), Some("hello"));
+        assert_eq!(buffer.push(&[0xE2, 0x82, 0xAC]).as_deref(), Some("€"));
+    }
+
+    #[test]
+    fn test_partial_utf8() {
+        let mut buffer = TokenUtf8Buffer::new();
+        assert_eq!(buffer.push(&[0xE2, 0x82]).as_deref(), None);
+        assert_eq!(buffer.push(&[0xAC]).as_deref(), Some("€"));
+    }
+
+    #[test]
+    fn test_invalid_prelude_for_valid_utf8() {
+        let mut buffer = TokenUtf8Buffer::new();
+        assert_eq!(buffer.push(&[0xD8]).as_deref(), None);
+        assert_eq!(buffer.push(&[0xE2, 0x82]).as_deref(), None);
+        assert_eq!(buffer.push(&[0xAC]).as_deref(), Some("€"));
+    }
+}