| 1 | from .configs import prompts |
| 2 | |
| 3 | class HtmlPrompt: |
| 4 | def __init__(self, prompt: str='') -> None: |
| 5 | prompt = self.extract(prompt, 'xml') |
| 6 | if prompt not in prompts: |
| 7 | raise Exception('Unknown prompt: ' + prompt) |
| 8 | |
| 9 | constructors = { |
| 10 | 'refine': self.normal_prompt_constructor, |
| 11 | 'xml': self.normal_prompt_constructor, |
| 12 | 'new_data': self.new_data_prompt_constructor, |
| 13 | } |
| 14 | |
| 15 | self.name = prompt |
| 16 | self.prompt = prompts[prompt] |
| 17 | self.constructor = constructors[prompt] |
| 18 | |
| 19 | @staticmethod |
| 20 | def extract(data, default=''): |
| 21 | return data if data is not None else default |
| 22 | |
| 23 | def subtree_constructor(self, subtree: list[str]=[]) -> str: |
| 24 | return self.prompt['subtree_splitter'].join(subtree) |
| 25 | |
| 26 | def normal_prompt_constructor(self, tag: str='', label: str='', content: str='', subtree_str: str='', class_dict: dict[str]={}) -> str: |
| 27 | def add_prefix(data, prefix): |
| 28 | return prefix + data if len(data) > 0 else '' |
| 29 | |
| 30 | tag = self.extract(tag) |
| 31 | label = self.extract(label) |
| 32 | content = self.extract(content) |
| 33 | subtree_str = self.extract(subtree_str, '') |
| 34 | class_dict = self.extract(class_dict, {}) |
| 35 | |
| 36 | label_str = '' |
| 37 | if len(label) > 0: |
| 38 | label_str = self.prompt['label'].format(label=label) |
| 39 | |
| 40 | classes = [] |
| 41 | values = set() |
| 42 | for key, val in class_dict.items(): |
| 43 | if val in values: |
| 44 | continue |
| 45 | values.add(val) |
| 46 | classes.append(self.prompt['attr'].format(key=key, attr=val)) |
| 47 | classes_str = self.prompt['attr_splitter'].join(classes) |
| 48 | |
| 49 | content_splitter = ' ' if len(classes_str) == 0 else self.prompt['attr_splitter'] |
| 50 | classes_str = add_prefix(classes_str, ' ') |
| 51 | content_str = add_prefix(content, content_splitter) |
| 52 | subtree_str = add_prefix(subtree_str, ' ') |
| 53 | |
| 54 | return self.prompt['dom'].format(tag=tag, label=label_str, attr=classes_str, content=content_str, subtree=subtree_str) |
| 55 | |
| 56 | def new_data_prompt_constructor(self, tag: str='', label: str='', content: str='', subtree_str: str='', class_dict: dict[str]={}) -> str: |
| 57 | def add_prefix(data, prefix): |
| 58 | return prefix + data if len(data) > 0 else '' |
| 59 | |
| 60 | tag = self.extract(tag) |