1use runmat_builtins::{
4 BuiltinCompletionPolicy, BuiltinDescriptor, BuiltinErrorDescriptor, BuiltinOutputMode,
5 BuiltinParamArity, BuiltinParamDescriptor, BuiltinParamType, BuiltinSignatureDescriptor,
6 CellArray, ObjectInstance, ResolveContext, StringArray, Type, Value,
7};
8use runmat_macros::runtime_builtin;
9
10use crate::builtins::strings::core::compat::scalar_text;
11use crate::builtins::strings::text_analytics::documents::{
12 document_token_type_with_options, documents_from_object, options_from_document_object,
13 replace_tokenized_document_documents, text_analytics_error, tokenized_document_language,
14 DocumentTokenType, TOKENIZED_DOCUMENT_CLASS,
15};
16use crate::builtins::strings::text_analytics::pos::{
17 add_part_of_speech_details_builtin, part_of_speech_details_from_object, POS_DETAILS_PROPERTY,
18};
19use crate::{gather_if_needed_async, BuiltinResult};
20
21pub(in crate::builtins::strings::text_analytics) const ENTITY_DETAILS_PROPERTY: &str =
22 "EntityDetails";
23
24const OUT_DOCUMENTS: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
25 name: "updatedDocuments",
26 ty: BuiltinParamType::Any,
27 arity: BuiltinParamArity::Required,
28 default: None,
29 description: "Updated tokenized document object.",
30}];
31
32const IN_DOCUMENTS: [BuiltinParamDescriptor; 1] = [BuiltinParamDescriptor {
33 name: "documents",
34 ty: BuiltinParamType::Any,
35 arity: BuiltinParamArity::Required,
36 default: None,
37 description: "tokenizedDocument object.",
38}];
39
40const IN_DOCUMENTS_REST: [BuiltinParamDescriptor; 2] = [
41 BuiltinParamDescriptor {
42 name: "documents",
43 ty: BuiltinParamType::Any,
44 arity: BuiltinParamArity::Required,
45 default: None,
46 description: "tokenizedDocument object.",
47 },
48 BuiltinParamDescriptor {
49 name: "NameValue",
50 ty: BuiltinParamType::Any,
51 arity: BuiltinParamArity::Variadic,
52 default: None,
53 description: "Name-value options: RetokenizeMethod, DiscardKnownValues, Model.",
54 },
55];
56
57const ERROR_INVALID_INPUT: BuiltinErrorDescriptor = BuiltinErrorDescriptor {
58 code: "RM.TEXT_ANALYTICS.ADD_ENTITY_DETAILS.INVALID_INPUT",
59 identifier: Some("RunMat:addEntityDetails:InvalidInput"),
60 when: "Input is not a supported tokenizedDocument object or option form.",
61 message: "addEntityDetails: invalid input",
62};
63
64const ERRORS: [BuiltinErrorDescriptor; 1] = [ERROR_INVALID_INPUT];
65
66pub const ADD_ENTITY_DETAILS_DESCRIPTOR: BuiltinDescriptor = BuiltinDescriptor {
67 signatures: &[
68 BuiltinSignatureDescriptor {
69 label: "updatedDocuments = addEntityDetails(documents)",
70 inputs: &IN_DOCUMENTS,
71 outputs: &OUT_DOCUMENTS,
72 },
73 BuiltinSignatureDescriptor {
74 label: "updatedDocuments = addEntityDetails(documents,Name,Value)",
75 inputs: &IN_DOCUMENTS_REST,
76 outputs: &OUT_DOCUMENTS,
77 },
78 ],
79 output_mode: BuiltinOutputMode::Fixed,
80 completion_policy: BuiltinCompletionPolicy::Public,
81 errors: &ERRORS,
82};
83
84fn any_type(_args: &[Type], _ctx: &ResolveContext) -> Type {
85 Type::Unknown
86}
87
88#[runtime_builtin(
89 name = "addEntityDetails",
90 category = "strings/text_analytics",
91 summary = "Add named-entity details to tokenizedDocument objects.",
92 keywords = "addEntityDetails,text analytics,tokenizedDocument,entity,named entity,ner",
93 accel = "sink",
94 type_resolver(any_type),
95 descriptor(crate::builtins::strings::text_analytics::entities::ADD_ENTITY_DETAILS_DESCRIPTOR),
96 builtin_path = "crate::builtins::strings::text_analytics::entities"
97)]
98async fn add_entity_details_builtin(args: Vec<Value>) -> BuiltinResult<Value> {
99 let gathered = gather_args(args).await?;
100 let (documents, options) = parse_args(gathered)?;
101 let mut object = tokenized_document_object(documents)?;
102 let language = EntityLanguage::from_document(&tokenized_document_language(&object))?;
103
104 if options.retokenize_method == RetokenizeMethod::Entity {
105 let documents = documents_from_object(&object, "addEntityDetails")?;
106 let retokenized = documents
107 .iter()
108 .map(|doc| retokenize_for_entities(doc, language))
109 .collect::<Vec<_>>();
110 if retokenized != documents {
111 replace_tokenized_document_documents(&mut object, retokenized, "addEntityDetails")?;
112 clear_token_aligned_details(&mut object);
113 }
114 }
115
116 if !object.properties.contains_key(POS_DETAILS_PROPERTY)
117 || !object.properties.contains_key("SentenceNumbers")
118 {
119 object = add_part_of_speech_details_for_entities(object).await?;
120 }
121
122 let documents = documents_from_object(&object, "addEntityDetails")?;
123 let stored = if options.discard_known_values {
124 None
125 } else {
126 entity_details_from_object(&object, "addEntityDetails")?
127 };
128 let document_options = options_from_document_object(&object);
129 let entity_tags = entity_details_for_documents(
130 &documents,
131 stored.as_deref(),
132 language,
133 &document_options,
134 options.discard_known_values,
135 )?;
136 object
137 .properties
138 .insert(ENTITY_DETAILS_PROPERTY.to_string(), entity_tags);
139 mark_entity_pos_as_proper_nouns(&mut object)?;
140 Ok(Value::Object(object))
141}
142
143async fn gather_args(args: Vec<Value>) -> BuiltinResult<Vec<Value>> {
144 let mut out = Vec::with_capacity(args.len());
145 for arg in args {
146 out.push(gather_if_needed_async(&arg).await.map_err(|err| {
147 text_analytics_error(
148 "addEntityDetails",
149 format!("addEntityDetails: failed to gather input: {err}"),
150 )
151 })?);
152 }
153 Ok(out)
154}
155
156#[derive(Clone, Debug)]
157struct AddEntityOptions {
158 discard_known_values: bool,
159 retokenize_method: RetokenizeMethod,
160}
161
162impl Default for AddEntityOptions {
163 fn default() -> Self {
164 Self {
165 discard_known_values: false,
166 retokenize_method: RetokenizeMethod::Entity,
167 }
168 }
169}
170
171#[derive(Clone, Copy, Debug, PartialEq, Eq)]
172enum RetokenizeMethod {
173 Entity,
174 None,
175}
176
177impl RetokenizeMethod {
178 fn parse(value: &str) -> BuiltinResult<Self> {
179 match value.trim().to_ascii_lowercase().as_str() {
180 "entity" => Ok(Self::Entity),
181 "none" => Ok(Self::None),
182 other => Err(text_analytics_error(
183 "addEntityDetails",
184 format!("addEntityDetails: unsupported RetokenizeMethod '{other}'"),
185 )),
186 }
187 }
188}
189
190#[derive(Clone, Copy, Debug, PartialEq, Eq)]
191enum EntityLanguage {
192 English,
193 Japanese,
194 German,
195 Korean,
196}
197
198impl EntityLanguage {
199 fn from_document(language: &str) -> BuiltinResult<Self> {
200 match language.trim().to_ascii_lowercase().as_str() {
201 "en" => Ok(Self::English),
202 "ja" => Ok(Self::Japanese),
203 "de" => Ok(Self::German),
204 "ko" => Ok(Self::Korean),
205 other => Err(text_analytics_error(
206 "addEntityDetails",
207 format!("addEntityDetails: unsupported document language '{other}'"),
208 )),
209 }
210 }
211}
212
213fn parse_args(args: Vec<Value>) -> BuiltinResult<(Value, AddEntityOptions)> {
214 if args.is_empty() {
215 return Err(text_analytics_error(
216 "addEntityDetails",
217 "addEntityDetails: expected tokenizedDocument input",
218 ));
219 }
220 if !(args.len() - 1).is_multiple_of(2) {
221 return Err(text_analytics_error(
222 "addEntityDetails",
223 "addEntityDetails: name-value options must appear in pairs",
224 ));
225 }
226 let mut options = AddEntityOptions::default();
227 let mut idx = 1usize;
228 while idx < args.len() {
229 let name = scalar_text(&args[idx], "addEntityDetails")
230 .map_err(|err| text_analytics_error("addEntityDetails", err.to_string()))?;
231 if name.eq_ignore_ascii_case("DiscardKnownValues") {
232 options.discard_known_values = logical_scalar(&args[idx + 1])?;
233 } else if name.eq_ignore_ascii_case("RetokenizeMethod") {
234 let value = scalar_text(&args[idx + 1], "addEntityDetails")
235 .map_err(|err| text_analytics_error("addEntityDetails", err.to_string()))?;
236 options.retokenize_method = RetokenizeMethod::parse(&value)?;
237 } else if name.eq_ignore_ascii_case("Model") {
238 parse_model_option(&args[idx + 1])?;
239 } else {
240 return Err(text_analytics_error(
241 "addEntityDetails",
242 format!("addEntityDetails: unsupported option '{name}'"),
243 ));
244 }
245 idx += 2;
246 }
247 Ok((args[0].clone(), options))
248}
249
250fn parse_model_option(value: &Value) -> BuiltinResult<()> {
251 if scalar_text(value, "addEntityDetails")
252 .map(|model| model.eq_ignore_ascii_case("auto"))
253 .unwrap_or(false)
254 {
255 return Ok(());
256 }
257 Err(text_analytics_error(
258 "addEntityDetails",
259 "addEntityDetails: only Model value 'auto' is supported; custom entity models remain tracked",
260 ))
261}
262
263fn tokenized_document_object(value: Value) -> BuiltinResult<ObjectInstance> {
264 match value {
265 Value::Object(object) if object.is_class(TOKENIZED_DOCUMENT_CLASS) => Ok(object),
266 Value::Object(object) => Err(text_analytics_error(
267 "addEntityDetails",
268 format!(
269 "addEntityDetails: expected tokenizedDocument object, got {}",
270 object.class_name
271 ),
272 )),
273 other => Err(text_analytics_error(
274 "addEntityDetails",
275 format!("addEntityDetails: expected tokenizedDocument object, got {other:?}"),
276 )),
277 }
278}
279
280fn clear_token_aligned_details(object: &mut ObjectInstance) {
281 for property in [
282 "TypeDetails",
283 "SentenceNumbers",
284 "LemmaDetails",
285 POS_DETAILS_PROPERTY,
286 ENTITY_DETAILS_PROPERTY,
287 "HeadDetails",
288 "DependencyDetails",
289 ] {
290 object.properties.remove(property);
291 }
292}
293
294async fn add_part_of_speech_details_for_entities(
295 object: ObjectInstance,
296) -> BuiltinResult<ObjectInstance> {
297 let Value::Object(object) = add_part_of_speech_details_builtin(vec![
298 Value::Object(object),
299 Value::String("RetokenizeMethod".to_string()),
300 Value::String("none".to_string()),
301 ])
302 .await?
303 else {
304 return Err(text_analytics_error(
305 "addEntityDetails",
306 "addEntityDetails: addPartOfSpeechDetails did not return tokenizedDocument",
307 ));
308 };
309 Ok(object)
310}
311
312fn entity_details_for_documents(
313 documents: &[Vec<String>],
314 stored: Option<&[Vec<String>]>,
315 language: EntityLanguage,
316 options: &crate::builtins::strings::text_analytics::documents::DocumentOptions,
317 discard_known_values: bool,
318) -> BuiltinResult<Value> {
319 let values = documents
320 .iter()
321 .enumerate()
322 .map(|(doc_idx, doc)| {
323 let tags = doc
324 .iter()
325 .enumerate()
326 .map(|(token_idx, token)| {
327 if !discard_known_values {
328 if let Some(tag) = stored
329 .and_then(|values| values.get(doc_idx))
330 .and_then(|values| values.get(token_idx))
331 .filter(|tag| is_known_entity(tag))
332 {
333 return tag.clone();
334 }
335 }
336 entity_for_token(token, language, options).to_string()
337 })
338 .collect::<Vec<_>>();
339 StringArray::new(tags, vec![1, doc.len()])
340 .map(Value::StringArray)
341 .map_err(|err| text_analytics_error("addEntityDetails", err))
342 })
343 .collect::<BuiltinResult<Vec<_>>>()?;
344 Ok(Value::Cell(
345 CellArray::new(values, documents.len(), 1)
346 .map_err(|err| text_analytics_error("addEntityDetails", err))?,
347 ))
348}
349
350fn entity_for_token(
351 token: &str,
352 language: EntityLanguage,
353 options: &crate::builtins::strings::text_analytics::documents::DocumentOptions,
354) -> &'static str {
355 if token.split_whitespace().count() > 1 {
356 return match language {
357 EntityLanguage::English => english_entity(token),
358 EntityLanguage::German => german_entity(token),
359 EntityLanguage::Japanese | EntityLanguage::Korean => asian_entity(token),
360 };
361 }
362 match document_token_type_with_options(token, options) {
363 DocumentTokenType::Punctuation | DocumentTokenType::Digits => "non-entity",
364 DocumentTokenType::WebAddress
365 | DocumentTokenType::EmailAddress
366 | DocumentTokenType::Hashtag
367 | DocumentTokenType::AtMention
368 | DocumentTokenType::Emoticon
369 | DocumentTokenType::Emoji
370 | DocumentTokenType::Other => "non-entity",
371 DocumentTokenType::Letters => match language {
372 EntityLanguage::English => english_entity(token),
373 EntityLanguage::German => german_entity(token),
374 EntityLanguage::Japanese | EntityLanguage::Korean => asian_entity(token),
375 },
376 }
377}
378
379fn english_entity(token: &str) -> &'static str {
380 let compact = compact_entity_token(token);
381 let lower = compact.to_ascii_lowercase();
382 if matches!(
383 lower.as_str(),
384 "mary"
385 | "john"
386 | "jane"
387 | "michael"
388 | "sarah"
389 | "david"
390 | "robert"
391 | "linda"
392 | "elizabeth"
393 | "william"
394 | "james"
395 ) {
396 return "person";
397 }
398 if matches!(
399 lower.as_str(),
400 "mathworks"
401 | "openai"
402 | "microsoft"
403 | "google"
404 | "apple"
405 | "amazon"
406 | "runmat"
407 | "volkswagen"
408 ) {
409 return "organization";
410 }
411 if matches!(
412 lower.as_str(),
413 "natick"
414 | "massachusetts"
415 | "boston"
416 | "california"
417 | "london"
418 | "paris"
419 | "tokyo"
420 | "seoul"
421 | "berlin"
422 | "munich"
423 | "newyork"
424 | "sanfrancisco"
425 | "unitedstates"
426 | "usa"
427 | "u.s.a."
428 ) {
429 return "location";
430 }
431 if token.chars().next().is_some_and(char::is_uppercase) {
432 "other"
433 } else {
434 "non-entity"
435 }
436}
437
438fn german_entity(token: &str) -> &'static str {
439 let compact = compact_entity_token(token);
440 let lower = compact.to_ascii_lowercase();
441 if matches!(
442 lower.as_str(),
443 "hans" | "anna" | "maria" | "peter" | "johann" | "max" | "sophie"
444 ) {
445 return "person";
446 }
447 if matches!(
448 lower.as_str(),
449 "volkswagen" | "sap" | "siemens" | "bmw" | "mercedes" | "mathworks" | "runmat"
450 ) {
451 return "organization";
452 }
453 if matches!(
454 lower.as_str(),
455 "berlin" | "münchen" | "munich" | "hamburg" | "wolfsburg" | "deutschland" | "germany"
456 ) {
457 return "location";
458 }
459 if token.chars().next().is_some_and(char::is_uppercase) {
460 "other"
461 } else {
462 "non-entity"
463 }
464}
465
466fn asian_entity(token: &str) -> &'static str {
467 let lower = token.to_ascii_lowercase();
468 if matches!(lower.as_str(), "tokyo" | "seoul" | "japan" | "korea") {
469 "location"
470 } else {
471 "non-entity"
472 }
473}
474
475fn compact_entity_token(token: &str) -> String {
476 token
477 .chars()
478 .filter(|ch| ch.is_alphanumeric() || *ch == '.')
479 .collect()
480}
481
482fn retokenize_for_entities(tokens: &[String], language: EntityLanguage) -> Vec<String> {
483 let mut out = Vec::new();
484 let mut idx = 0usize;
485 while idx < tokens.len() {
486 if let Some((merged, consumed)) = multiword_entity_at(tokens, idx, language) {
487 out.push(merged);
488 idx += consumed;
489 } else {
490 out.push(tokens[idx].clone());
491 idx += 1;
492 }
493 }
494 out
495}
496
497fn multiword_entity_at(
498 tokens: &[String],
499 idx: usize,
500 language: EntityLanguage,
501) -> Option<(String, usize)> {
502 let token = tokens.get(idx)?;
503 let mut parts = vec![token.clone()];
504 let mut cursor = idx + 1;
505 while let Some(next) = tokens.get(cursor) {
506 if next == "." || next == "," || next == "-" {
507 break;
508 }
509 if !looks_like_entity_word(next, language) {
510 break;
511 }
512 parts.push(next.clone());
513 cursor += 1;
514 }
515 if parts.len() < 2 || !looks_like_entity_word(token, language) {
516 return None;
517 }
518 let merged = parts.join(" ");
519 let tag = match language {
520 EntityLanguage::English => english_entity(&merged),
521 EntityLanguage::German => german_entity(&merged),
522 EntityLanguage::Japanese | EntityLanguage::Korean => asian_entity(&merged),
523 };
524 if tag == "non-entity" {
525 None
526 } else {
527 Some((merged, cursor - idx))
528 }
529}
530
531fn looks_like_entity_word(token: &str, language: EntityLanguage) -> bool {
532 if !token.chars().next().is_some_and(char::is_uppercase) {
533 return false;
534 }
535 match language {
536 EntityLanguage::English => english_entity(token) != "non-entity",
537 EntityLanguage::German => german_entity(token) != "non-entity",
538 EntityLanguage::Japanese | EntityLanguage::Korean => false,
539 }
540}
541
542fn mark_entity_pos_as_proper_nouns(object: &mut ObjectInstance) -> BuiltinResult<()> {
543 let Some(entities) = entity_details_from_object(object, "addEntityDetails")? else {
544 return Ok(());
545 };
546 let Some(mut pos) = part_of_speech_details_from_object(object, "addEntityDetails")? else {
547 return Ok(());
548 };
549 for (doc_idx, entity_doc) in entities.iter().enumerate() {
550 if let Some(pos_doc) = pos.get_mut(doc_idx) {
551 for (token_idx, entity) in entity_doc.iter().enumerate() {
552 if entity != "non-entity" {
553 if let Some(pos_tag) = pos_doc.get_mut(token_idx) {
554 *pos_tag = "proper-noun".to_string();
555 }
556 }
557 }
558 }
559 }
560 let values = pos
561 .iter()
562 .map(|doc| {
563 StringArray::new(doc.clone(), vec![1, doc.len()])
564 .map(Value::StringArray)
565 .map_err(|err| text_analytics_error("addEntityDetails", err))
566 })
567 .collect::<BuiltinResult<Vec<_>>>()?;
568 object.properties.insert(
569 POS_DETAILS_PROPERTY.to_string(),
570 Value::Cell(
571 CellArray::new(values, pos.len(), 1)
572 .map_err(|err| text_analytics_error("addEntityDetails", err))?,
573 ),
574 );
575 Ok(())
576}
577
578fn is_known_entity(value: &str) -> bool {
579 let trimmed = value.trim();
580 !trimmed.is_empty()
581 && trimmed != "unknown"
582 && !crate::builtins::strings::common::is_missing_string(trimmed)
583}
584
585pub(in crate::builtins::strings::text_analytics) fn entity_details_from_object(
586 object: &ObjectInstance,
587 fn_name: &str,
588) -> BuiltinResult<Option<Vec<Vec<String>>>> {
589 let Some(value) = object.properties.get(ENTITY_DETAILS_PROPERTY) else {
590 return Ok(None);
591 };
592 let Value::Cell(cell) = value else {
593 return Err(text_analytics_error(
594 fn_name,
595 format!("{fn_name}: tokenizedDocument object has invalid EntityDetails property"),
596 ));
597 };
598 if cell.cols != 1 {
599 return Err(text_analytics_error(
600 fn_name,
601 format!("{fn_name}: tokenizedDocument object has invalid EntityDetails shape"),
602 ));
603 }
604 let mut out = Vec::with_capacity(cell.data.len());
605 for item in &cell.data {
606 let Value::StringArray(array) = item else {
607 return Err(text_analytics_error(
608 fn_name,
609 format!("{fn_name}: tokenizedDocument object has invalid EntityDetails entry"),
610 ));
611 };
612 if array.rows != 1 {
613 return Err(text_analytics_error(
614 fn_name,
615 format!(
616 "{fn_name}: tokenizedDocument object has invalid EntityDetails entry shape"
617 ),
618 ));
619 }
620 out.push(array.data.clone());
621 }
622 Ok(Some(out))
623}
624
625fn logical_scalar(value: &Value) -> BuiltinResult<bool> {
626 match value {
627 Value::Bool(value) => Ok(*value),
628 Value::Num(value) if *value == 0.0 || *value == 1.0 => Ok(*value != 0.0),
629 Value::Tensor(tensor) if tensor.data.len() == 1 => match tensor.data[0] {
630 0.0 => Ok(false),
631 1.0 => Ok(true),
632 other => Err(text_analytics_error(
633 "addEntityDetails",
634 format!(
635 "addEntityDetails: logical scalar option must be true or false, got {other}"
636 ),
637 )),
638 },
639 Value::LogicalArray(array) if array.data.len() == 1 => Ok(array.data[0] != 0),
640 other => Err(text_analytics_error(
641 "addEntityDetails",
642 format!("addEntityDetails: logical scalar option must be true or false, got {other:?}"),
643 )),
644 }
645}
646
647#[cfg(test)]
648mod tests {
649 use super::*;
650 use crate::builtins::strings::text_analytics::details::token_details_builtin;
651 use crate::builtins::strings::text_analytics::documents::tokenized_document_builtin;
652 use crate::builtins::table::{table_variable_names_from_object, table_variables};
653 use runmat_builtins::{LogicalArray, Tensor};
654
655 fn run_tokenized(args: Vec<Value>) -> BuiltinResult<Value> {
656 futures::executor::block_on(tokenized_document_builtin(args))
657 }
658
659 fn run_add_entity(args: Vec<Value>) -> BuiltinResult<Value> {
660 futures::executor::block_on(add_entity_details_builtin(args))
661 }
662
663 fn run_token_details(value: Value) -> BuiltinResult<Value> {
664 futures::executor::block_on(token_details_builtin(value))
665 }
666
667 fn object(value: Value) -> ObjectInstance {
668 let Value::Object(object) = value else {
669 panic!("expected object");
670 };
671 object
672 }
673
674 fn table_column(table: &ObjectInstance, name: &str) -> Value {
675 table_variables(table)
676 .expect("table variables")
677 .fields
678 .get(name)
679 .cloned()
680 .unwrap_or_else(|| panic!("missing table column {name}"))
681 }
682
683 fn string_column(table: &ObjectInstance, name: &str) -> Vec<String> {
684 match table_column(table, name) {
685 Value::StringArray(array) => array.data,
686 other => panic!("expected string column {name}, got {other:?}"),
687 }
688 }
689
690 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
691 #[test]
692 fn add_entity_details_adds_sentence_pos_and_entity_columns() {
693 let docs =
694 run_tokenized(vec![Value::String("Mary uses MATLAB at MathWorks.".into())]).unwrap();
695 let updated = run_add_entity(vec![docs]).expect("entities");
696 let table = object(run_token_details(updated).expect("details"));
697 assert_eq!(
698 table_variable_names_from_object(&table).expect("names"),
699 vec![
700 "Token",
701 "DocumentNumber",
702 "SentenceNumber",
703 "LineNumber",
704 "Type",
705 "Language",
706 "PartOfSpeech",
707 "Entity"
708 ]
709 );
710 assert_eq!(
711 string_column(&table, "Entity"),
712 vec![
713 "person",
714 "non-entity",
715 "other",
716 "non-entity",
717 "organization",
718 "non-entity"
719 ]
720 );
721 assert_eq!(
722 string_column(&table, "PartOfSpeech"),
723 vec![
724 "proper-noun",
725 "verb",
726 "proper-noun",
727 "adposition",
728 "proper-noun",
729 "punctuation"
730 ]
731 );
732 }
733
734 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
735 #[test]
736 fn add_entity_details_retokenizes_multiword_entities() {
737 let docs = run_tokenized(vec![Value::String("Mary moved to New York.".into())]).unwrap();
738 let updated = run_add_entity(vec![docs]).expect("entities");
739 let table = object(run_token_details(updated).expect("details"));
740 assert_eq!(
741 string_column(&table, "Token"),
742 vec!["Mary", "moved", "to", "New York", "."]
743 );
744 assert_eq!(
745 string_column(&table, "Entity"),
746 vec![
747 "person",
748 "non-entity",
749 "non-entity",
750 "location",
751 "non-entity"
752 ]
753 );
754 }
755
756 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
757 #[test]
758 fn add_entity_details_preserves_existing_known_values_unless_discarding() {
759 let docs = object(
760 run_tokenized(vec![
761 Value::StringArray(
762 StringArray::new(vec!["dogs".into(), "run".into()], vec![1, 2]).unwrap(),
763 ),
764 Value::String("TokenizeMethod".into()),
765 Value::String("none".into()),
766 ])
767 .expect("tokenized"),
768 );
769 let mut stale = docs.clone();
770 stale.properties.insert(
771 ENTITY_DETAILS_PROPERTY.to_string(),
772 Value::Cell(
773 CellArray::new(
774 vec![Value::StringArray(
775 StringArray::new(vec!["custom".into(), "".into()], vec![1, 2]).unwrap(),
776 )],
777 1,
778 1,
779 )
780 .unwrap(),
781 ),
782 );
783
784 let preserved = run_add_entity(vec![
785 Value::Object(stale.clone()),
786 Value::String("RetokenizeMethod".into()),
787 Value::String("none".into()),
788 ])
789 .expect("preserve");
790 let table = object(run_token_details(preserved).expect("details"));
791 assert_eq!(
792 string_column(&table, "Entity"),
793 vec!["custom", "non-entity"]
794 );
795
796 let recomputed = run_add_entity(vec![
797 Value::Object(stale),
798 Value::String("RetokenizeMethod".into()),
799 Value::String("none".into()),
800 Value::String("DiscardKnownValues".into()),
801 Value::LogicalArray(LogicalArray::new(vec![1], vec![1, 1]).unwrap()),
802 ])
803 .expect("recompute");
804 let table = object(run_token_details(recomputed).expect("details"));
805 assert_eq!(
806 string_column(&table, "Entity"),
807 vec!["non-entity", "non-entity"]
808 );
809 }
810
811 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
812 #[test]
813 fn add_entity_details_adds_sentence_details_when_pos_already_exists() {
814 let docs = object(
815 run_tokenized(vec![
816 Value::StringArray(
817 StringArray::new(vec!["Mary".into(), "runs".into()], vec![1, 2]).unwrap(),
818 ),
819 Value::String("TokenizeMethod".into()),
820 Value::String("none".into()),
821 ])
822 .expect("tokenized"),
823 );
824 let mut with_pos = docs;
825 with_pos.properties.insert(
826 POS_DETAILS_PROPERTY.to_string(),
827 Value::Cell(
828 CellArray::new(
829 vec![Value::StringArray(
830 StringArray::new(vec!["custom".into(), "verb".into()], vec![1, 2]).unwrap(),
831 )],
832 1,
833 1,
834 )
835 .unwrap(),
836 ),
837 );
838
839 let updated = run_add_entity(vec![
840 Value::Object(with_pos),
841 Value::String("RetokenizeMethod".into()),
842 Value::String("none".into()),
843 ])
844 .expect("entities");
845 let table = object(run_token_details(updated).expect("details"));
846 assert!(table_variable_names_from_object(&table)
847 .expect("names")
848 .contains(&"SentenceNumber".to_string()));
849 assert_eq!(
850 string_column(&table, "PartOfSpeech"),
851 vec!["proper-noun", "verb"]
852 );
853 }
854
855 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
856 #[test]
857 fn add_entity_details_rejects_bad_options_languages_and_models() {
858 let docs = run_tokenized(vec![Value::String("Mary runs".into())]).unwrap();
859 let err = run_add_entity(vec![
860 docs.clone(),
861 Value::String("RetokenizeMethod".into()),
862 Value::String("bogus".into()),
863 ])
864 .expect_err("bad retokenizer");
865 assert!(err.to_string().contains("unsupported RetokenizeMethod"));
866
867 let err = run_add_entity(vec![
868 docs.clone(),
869 Value::String("Model".into()),
870 Value::String("custom".into()),
871 ])
872 .expect_err("bad model");
873 assert!(err.to_string().contains("only Model value"));
874
875 let mut unsupported = object(docs);
876 unsupported
877 .properties
878 .insert("Language".into(), Value::String("fr".into()));
879 let err =
880 run_add_entity(vec![Value::Object(unsupported)]).expect_err("unsupported language");
881 assert!(err.to_string().contains("unsupported document language"));
882 }
883
884 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
885 #[test]
886 fn token_details_rejects_malformed_entity_details() {
887 let docs = object(run_tokenized(vec![Value::String("Mary runs".into())]).unwrap());
888 let mut malformed = docs.clone();
889 malformed.properties.insert(
890 ENTITY_DETAILS_PROPERTY.to_string(),
891 Value::Cell(
892 CellArray::new(
893 vec![Value::StringArray(
894 StringArray::new(vec!["person".into()], vec![1, 1]).unwrap(),
895 )],
896 1,
897 1,
898 )
899 .unwrap(),
900 ),
901 );
902 let err = run_token_details(Value::Object(malformed)).expect_err("expected error");
903 assert!(err.to_string().contains("EntityDetails entry"));
904 }
905
906 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
907 #[test]
908 fn add_entity_details_discard_known_values_ignores_malformed_stored_tags() {
909 let docs = object(run_tokenized(vec![Value::String("Mary runs".into())]).unwrap());
910 let mut malformed = docs.clone();
911 malformed.properties.insert(
912 ENTITY_DETAILS_PROPERTY.to_string(),
913 Value::String("bad".into()),
914 );
915 let updated = run_add_entity(vec![
916 Value::Object(malformed),
917 Value::String("DiscardKnownValues".into()),
918 Value::Bool(true),
919 ])
920 .expect("discard malformed");
921 let table = object(run_token_details(updated).expect("details"));
922 assert_eq!(
923 string_column(&table, "Entity"),
924 vec!["person", "non-entity"]
925 );
926 }
927
928 #[cfg_attr(target_arch = "wasm32", wasm_bindgen_test::wasm_bindgen_test)]
929 #[test]
930 fn add_entity_details_accepts_german_language_metadata() {
931 let docs = object(
932 run_tokenized(vec![Value::String("Volkswagen ist in Wolfsburg.".into())]).unwrap(),
933 );
934 let mut german = docs;
935 german
936 .properties
937 .insert("Language".into(), Value::String("de".into()));
938 let updated = run_add_entity(vec![Value::Object(german)]).expect("german entities");
939 let table = object(run_token_details(updated).expect("details"));
940 assert_eq!(
941 string_column(&table, "Entity"),
942 vec![
943 "organization",
944 "non-entity",
945 "non-entity",
946 "location",
947 "non-entity"
948 ]
949 );
950 }
951
952 #[test]
953 fn logical_scalar_rejects_non_logical_numeric_values() {
954 let err = logical_scalar(&Value::Tensor(Tensor::new(vec![2.0], vec![1, 1]).unwrap()))
955 .expect_err("bad logical");
956 assert!(err.to_string().contains("logical scalar"));
957 }
958}