guillaume-be/rust-bert
1
1
Fork 0
mirror of https://github.com/guillaume-be/rust-bert.git synced 2024-08-16 16:10:25 +03:00
Code Issues Packages Projects Releases Wiki Activity

Generalize input of NERModel::predict_full_entities (#329)

Browse Source
This commit is contained in:
Åke Amcoff 2023-01-23 21:24:18 +01:00 committed by GitHub
parent e1e8fc615d
commit 80e0197e2c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
src/pipelines/ner.rs
View File

@ -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();