Fix mute tokens dead byte

Author: agourdel

benchmarks/bench_regex_guide.py

@@ -21,7 +21,7 @@ class RegexIndexBenchmark:
     params = regex_samples.keys()
 
     def setup(self, pattern_name):
-        self.vocabulary = Vocabulary.from_pretrained("gpt2")
+        self.vocabulary = Vocabulary.from_pretrained("unsloth/Llama-3.1-8B-Instruct")
         self.pattern = regex_samples[pattern_name]
 
     def time_regex_to_guide(self, pattern_name):

benchmarks/test_index_time.py

@@ -0,0 +1,17 @@
+import timeit
+from outlines_core import Index, Vocabulary
+
+regex_samples = {
+    "email": r"[a-z0-9!#$%&'*+/=?^_`{|}~-]+(?:\.[a-z0-9!#$%&'*+/=?^_`{|}~-]+)*@(?:[a-z0-9](?:[a-z0-9-]*[a-z0-9])?\.)+[a-z0-9](?:[a-z0-9-]*[a-z0-9])?",
+    # Ajoute d'autres regex si nécessaire
+}
+
+# Initialisation du Vocabulary avant la mesure
+vocabulary = Vocabulary.from_pretrained("unsloth/Llama-3.1-8B-Instruct")
+pattern = regex_samples["email"]
+# Code de setup (ne contient que l'importation et la définition de pattern)
+setup_code = "from outlines_core import Index"
+# Mesure uniquement la construction de l'Index
+stmt = "Index(pattern, vocabulary)"
+execution_time = timeit.timeit(stmt, setup=setup_code, globals=locals(), number=1)
+print(f"Temps d'exécution pour une construction froide de l'Index (Vocabulary pré-initialisé) : {execution_time} secondes")

setup.py

@@ -13,6 +13,7 @@
         binding=Binding.PyO3,
         features=["python-bindings"],
         rustc_flags=["--crate-type=cdylib"],
+        debug=False,
     ),
 ]
 

src/tokens_dfa/mod.rs

@@ -57,6 +57,7 @@ impl TokensDFA
         //let start_muting = Instant::now();
         let (muted_regex, muted_list) = mute_literals(regex, vocabulary, &mut additionnal_tokens);
         //let time_muting = start_muting.elapsed();
+        // println!("> Muted Regex : {}", muted_regex);
         // println!("> Muted List {:?}", muted_list);
         // println!("> Additionnal : {:?}", additionnal_tokens);
         let alphabet_len = vocabulary.len_alphabet()+additionnal_tokens.len(); // Real number of different token_id

src/tokens_dfa/reduce.rs

@@ -38,7 +38,7 @@ pub fn init_classes_and_graph_optimized(
         if token_ids[0] == 216 {return None;} // BUG IN THE VOCABULARY. token_id 216 ""\u011c"" is interpreted as \x1C, (Byte 28)
         let t_class = get_token_class(token, byte_classes);
 
-        if t_class.as_bytes().iter().any(|byte| dead_byte_classes.contains(byte)) {
+        if t_class.as_bytes().iter().any(|byte| dead_byte_classes.contains(byte)) && !token_ids.iter().any(|&id| additionnal_tokens.iter().any(|(_, add_id)| *add_id == id)) {
             return None; 
         }
 
@@ -67,7 +67,6 @@ pub fn init_classes_and_graph_optimized(
             transitions_table.get_token_ids_by_class().add_token_id(class_id, id);
             // Avoid to override the Tokens class of muted token.
             if additionnal_classes.contains(&t_class){
-                
                 continue;
 
             }
@@ -313,9 +312,7 @@ fn decompose_all_literals_optimized(
 
             token_sequence.reverse();
             result.insert(literal.clone(), (token_sequence, positions.clone()));
-        } else {
-            println!("Aucune décomposition trouvée pour le littéral: {}", literal);
-        }
+        } 
     }
 
     result

src/tokens_dfa/transitions_table.rs

@@ -164,7 +164,6 @@ impl MasksTable {
                 let mut real_class_id = class_id.clone();
 
                 if tokens.len() == 1 && muted_list.contains(&tokens[0]) { 
-                    
                     real_class_id = self.token_classes[tokens[0]]
                 }
 

src/v2_index.rs

@@ -159,10 +159,10 @@ mod tests {
 
     #[test]
     fn test_sample(){
-        let regex = r"[a-z0-9!#$%&'*+/=?^_`{|}~-]+(?:\.[a-z0-9!#$%&'*+/=?^_`{|}~-]+)*@(?:[a-z0-9](?:[a-z0-9-]*[a-z0-9])?\.)+[a-z0-9](?:[a-z0-9-]*[a-z0-9])?";
+        let regex = r"(0?[1-9]|1[0-2]):[0-5]\d\s?(am|pm)?";
         //let sch =r#"{"type": "object", "properties": {"name": {"type": "string"}, "age": {"type": "integer"}, "complexe_phone": {"type": "string", "pattern": "\\+?\\d{1,4}?[-. ]?\\(\\d{1,3}\\)?[-. ]?\\d{1,4}[-. ]?\\d{1,4}[-. ]?\\d{1,9}"}}, "required": ["name", "age", "complexe_phone"]}"#;
         //let regex = &json_schema::regex_from_str(sch, None).unwrap(); 
-         println!("{}", regex);
+        println!("{}", regex);
         let model_name = "unsloth/Llama-3.1-8B-Instruct";
         let vocab = Vocabulary::from_pretrained(model_name, None).unwrap();