"""CLI tool for sampling sentences from intents.""" import argparse import itertools import json import logging import sys from functools import partial from pathlib import Path from typing import Dict, Iterable, Optional, Set, Tuple import yaml from unicode_rbnf import RbnfEngine from .errors import MissingListError, MissingRuleError from .expression import ( Alternative, Expression, Group, ListReference, Permutation, RuleReference, Sentence, Sequence, TextChunk, ) from .intents import Intents, RangeSlotList, SlotList, TextSlotList, WildcardSlotList from .util import merge_dict, normalize_whitespace _LOGGER = logging.getLogger("hassil.sample") # lang -> engine _ENGINE_CACHE: Dict[str, RbnfEngine] = {} def sample_intents( intents: Intents, slot_lists: Optional[Dict[str, SlotList]] = None, expansion_rules: Optional[Dict[str, Sentence]] = None, max_sentences_per_intent: Optional[int] = None, intent_names: Optional[Set[str]] = None, language: Optional[str] = None, exclude_sentences_with_wildcards: bool = True, expand_ranges: bool = True, ) -> Iterable[Tuple[str, str]]: """Sample text strings for sentences from intents.""" if slot_lists is None: slot_lists = intents.slot_lists else: # Combine with intents slot_lists = {**intents.slot_lists, **slot_lists} if slot_lists is None: slot_lists = {} if expansion_rules is None: expansion_rules = intents.expansion_rules else: # Combine rules expansion_rules = {**intents.expansion_rules, **expansion_rules} for intent_name, intent in intents.intents.items(): if intent_names and (intent_name not in intent_names): # Skip intent continue num_intent_sentences = 0 skip_intent = False for intent_data in intent.data: if intent_data.expansion_rules: local_expansion_rules = { **expansion_rules, **intent_data.expansion_rules, } else: local_expansion_rules = expansion_rules for intent_sentence in intent_data.sentences: if exclude_sentences_with_wildcards and any( list_name in intent_data.wildcard_list_names for list_name in intent_sentence.list_names(local_expansion_rules) ): continue sentence_texts = sample_sentence( intent_sentence, slot_lists, local_expansion_rules, language=language, expand_ranges=expand_ranges, ) for sentence_text in sentence_texts: yield (intent_name, sentence_text) num_intent_sentences += 1 if (max_sentences_per_intent is not None) and ( 0 < max_sentences_per_intent <= num_intent_sentences ): skip_intent = True break if skip_intent: break if skip_intent: break def sample_sentence( sentence: Sentence, slot_lists: Optional[Dict[str, SlotList]] = None, expansion_rules: Optional[Dict[str, Sentence]] = None, language: Optional[str] = None, expand_lists: bool = True, expand_ranges: bool = True, skip_optionals: bool = False, ) -> Iterable[str]: return sample_expression( sentence.expression, slot_lists, expansion_rules, language, expand_lists, expand_ranges, skip_optionals, ) def sample_expression( expression: Expression, slot_lists: Optional[Dict[str, SlotList]] = None, expansion_rules: Optional[Dict[str, Sentence]] = None, language: Optional[str] = None, expand_lists: bool = True, expand_ranges: bool = True, skip_optionals: bool = False, ) -> Iterable[str]: """Sample possible text strings from an expression.""" if isinstance(expression, TextChunk): chunk: TextChunk = expression yield chunk.original_text elif isinstance(expression, Group): grp: Group = expression if isinstance(grp, Alternative): alt: Alternative = grp if skip_optionals and alt.is_optional: yield "" else: for item in grp.items: yield from sample_expression( item, slot_lists, expansion_rules, language=language, expand_lists=expand_lists, expand_ranges=expand_ranges, skip_optionals=skip_optionals, ) elif isinstance(grp, Sequence): seq_sentences = map( partial( sample_expression, slot_lists=slot_lists, expansion_rules=expansion_rules, language=language, expand_lists=expand_lists, expand_ranges=expand_ranges, skip_optionals=skip_optionals, ), grp.items, ) sentence_texts = itertools.product(*seq_sentences) for sentence_words in sentence_texts: yield normalize_whitespace("".join(sentence_words)) elif isinstance(grp, Permutation): # Need to make lists instead because itertools does multiple passes. grp_sentences = [ list( sample_expression( item, slot_lists, expansion_rules, language=language, expand_lists=expand_lists, expand_ranges=expand_ranges, skip_optionals=skip_optionals, ) ) for item in grp.items ] for perm_sentences in itertools.permutations(grp_sentences): sentence_texts = itertools.product(*perm_sentences) for sentence_words in sentence_texts: # Strip added whitespace yield normalize_whitespace("".join(sentence_words)).strip() else: raise ValueError(f"Unexpected group type: {grp}") elif isinstance(expression, ListReference): # {list} list_ref: ListReference = expression if not expand_lists: yield f"{{{list_ref.list_name}}}" return if (not slot_lists) or (list_ref.list_name not in slot_lists): raise MissingListError(f"Missing slot list {{{list_ref.list_name}}}") slot_list = slot_lists[list_ref.list_name] if isinstance(slot_list, TextSlotList): text_list: TextSlotList = slot_list if not text_list.values: # Not necessarily an error, but may be a surprise _LOGGER.warning("No values for list: %s", list_ref.list_name) for text_value in text_list.values: yield from sample_expression( text_value.text_in, slot_lists, expansion_rules, language=language, expand_lists=expand_lists, expand_ranges=expand_ranges, skip_optionals=skip_optionals, ) elif isinstance(slot_list, RangeSlotList): range_list: RangeSlotList = slot_list if not expand_ranges: if range_list.name: yield f"{{{range_list.name}}}" else: yield "{number}" return if range_list.digits: number_strs = map(str, range_list.get_numbers()) yield from number_strs if range_list.words: words_language = range_list.words_language or language if words_language: engine = _ENGINE_CACHE.get(words_language) if engine is None: engine = RbnfEngine.for_language(words_language) _ENGINE_CACHE[words_language] = engine assert engine is not None # digits -> words for word_number in range_list.get_numbers(): # Use all unique words for a number, including different # genders, cases, etc. format_result = engine.format_number(word_number) unique_number_strs = set(format_result.text_by_ruleset.values()) yield from unique_number_strs else: _LOGGER.warning( "No language set, so cannot convert %s digits to words", list_ref.slot_name, ) elif isinstance(slot_list, WildcardSlotList): wildcard_list: WildcardSlotList = slot_list if wildcard_list.name: yield f"{{{wildcard_list.name}}}" else: yield "{wildcard}" else: raise ValueError(f"Unexpected slot list type: {slot_list}") elif isinstance(expression, RuleReference): # rule_ref: RuleReference = expression if (not expansion_rules) or (rule_ref.rule_name not in expansion_rules): raise MissingRuleError(f"Missing expansion rule <{rule_ref.rule_name}>") rule_body = expansion_rules[rule_ref.rule_name].expression yield from sample_expression( rule_body, slot_lists, expansion_rules, language=language, expand_lists=expand_lists, expand_ranges=expand_ranges, skip_optionals=skip_optionals, ) else: raise ValueError(f"Unexpected expression: {expression}") def main(): """Main entry point""" parser = argparse.ArgumentParser() parser.add_argument("yaml", nargs="+", help="YAML files or directories") parser.add_argument( "-n", "--max-sentences-per-intent", type=int, help="Limit number of sentences per intent", ) parser.add_argument( "--intents", nargs="+", help="Only sample sentences from these intents" ) parser.add_argument( "--areas", nargs="+", help="Area names", default=["area"], ) parser.add_argument( "--names", nargs="+", default=["entity"], help="Device/entity names" ) parser.add_argument("--language", help="Language for digits to words") parser.add_argument( "--debug", action="store_true", help="Print DEBUG messages to the console" ) args = parser.parse_args() level = logging.DEBUG if args.debug else logging.INFO logging.basicConfig(level=level) _LOGGER.debug(args) slot_lists = { "area": TextSlotList.from_strings(args.areas), "name": TextSlotList.from_strings(args.names), } input_dict = {"intents": {}} for yaml_path_str in args.yaml: yaml_path = Path(yaml_path_str) if yaml_path.is_dir(): yaml_file_paths = yaml_path.glob("*.yaml") else: yaml_file_paths = [yaml_path] for yaml_file_path in yaml_file_paths: _LOGGER.debug("Loading file: %s", yaml_file_path) with open(yaml_file_path, "r", encoding="utf-8") as yaml_file: merge_dict(input_dict, yaml.safe_load(yaml_file)) assert input_dict, "No intent YAML files loaded" intents = Intents.from_dict(input_dict) intents_and_texts = sample_intents( intents, slot_lists, max_sentences_per_intent=args.max_sentences_per_intent, intent_names=set(args.intents) if args.intents else None, language=args.language, ) for intent_name, sentence_text in intents_and_texts: json.dump( {"intent": intent_name, "text": sentence_text.strip()}, sys.stdout, ensure_ascii=False, ) print("") if __name__ == "__main__": main()