Generalize input of NERModel::predict_full_entities (guillaume-be#329)

SpirosMakris · Jan 23, 2023 · 80e0197 · 80e0197
1 parent e1e8fc6
commit 80e0197
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/src/pipelines/ner.rs b/src/pipelines/ner.rs
@@ -274,7 +274,10 @@ impl NERModel {
     /// ]]
     /// # ;
     /// ```
-    pub fn predict_full_entities(&self, input: &[&str]) -> Vec<Vec<Entity>> {
+    pub fn predict_full_entities<S>(&self, input: &[S]) -> Vec<Vec<Entity>>
+    where
+        S: AsRef<str>,
+    {
         let tokens = self.token_classification_model.predict(input, true, false);
         let mut entities: Vec<Vec<Entity>> = Vec::new();