use crate::Result;
use crate::types::entity::{Entity, EntityCategory};
use async_trait::async_trait;
#[cfg_attr(not(target_arch = "wasm32"), async_trait)]
#[cfg_attr(target_arch = "wasm32", async_trait(?Send))]
#[cfg_attr(alef, alef(skip))]
pub trait NerBackend: Send + Sync {
async fn detect(&self, text: &str, categories: &[EntityCategory]) -> Result<Vec<Entity>>;
async fn detect_with_custom(
&self,
text: &str,
categories: &[EntityCategory],
custom_labels: &[String],
) -> Result<Vec<Entity>> {
if custom_labels.is_empty() {
return self.detect(text, categories).await;
}
let mut all: Vec<EntityCategory> = categories.to_vec();
for label in custom_labels {
all.push(EntityCategory::Custom(label.clone()));
}
self.detect(text, &all).await
}
}