from typing import Dict, List, Tuple, Union
import re
from transformers import (
    CamembertTokenizer,
    AutoTokenizer,
    pipeline,
)


class ThaiNameTagger:
    def __init__(self,
                dataset_name: str = "thainer",
                model_name: str = "wangchanberta-base-att-spm-uncased"):
        self.model_name = model_name
        self.tokenizer = CamembertTokenizer.from_pretrained(
                                    f'airesearch/{self.model_name}',
                                    revision='main')
        if self.model_name == "wangchanberta-base-att-spm-uncased":
            self.tokenizer.additional_special_tokens = ['<s>NOTUSED', '</s>NOTUSED', '<_>']
        self.classify_tokens = pipeline(
            task='ner',
            tokenizer=self.tokenizer,
            model = f'airesearch/{self.model_name}',
            revision = f'finetuned@{dataset_name}-ner',
            ignore_labels=[], 
            grouped_entities=True
        )
    
    def IOB(self, tag):
        if tag != "O":
            return "B-"+tag
        return "O"

    def get_ner(
        self, text: str, tag: bool = False
    ) -> List[Tuple[str, str]]:
        """
        This function tags named-entitiy from text in IOB format.

        Powered by wangchanberta from VISTEC-depa AI Research Institute of Thailand
        :param str text: text in Thai to be tagged
        :param bool tag: output like html tag.
        :return: a list of tuple associated with tokenized word group, NER tag,
                 and output like html tag (if the parameter `tag` is
                 specified as `True`).
                 Otherwise, return a list of tuple associated with tokenized
                 word and NER tag
        :rtype: Union[list[tuple[str, str]]], str
        """
        text = re.sub(" ", "<_>", text)
        self.json_ner = self.classify_tokens(text)
        self.output = ""
        self.sent_ner = [(i['word'].replace("<_>", " "),self.IOB(i['entity_group'])) for i in self.json_ner]
        if tag:
            temp = ""
            sent = ""
            for idx, (word, ner) in enumerate(self.sent_ner):
                if ner.startswith("B-") and temp != "":
                    sent += "</" + temp + ">"
                    temp = ner[2:]
                    sent += "<" + temp + ">"
                elif ner.startswith("B-"):
                    temp = ner[2:]
                    sent += "<" + temp + ">"
                elif ner == "O" and temp != "":
                    sent += "</" + temp + ">"
                    temp = ""
                sent += word

                if idx == len(self.sent_ner) - 1 and temp != "":
                    sent += "</" + temp + ">"

            return sent
        
        return self.sent_ner

@@ -301,47 +301,184 @@
 sed%22
-):%0A        self.model_name = model_name
+,%0A                grouped_entities: bool = True):%0A        self.model_name = model_name%0A        self.dataset_name = dataset_name%0A        self.grouped_entities = grouped_entities
 %0A   
@@ -1009,16 +1009,21 @@
 etuned@%7B
+self.
 dataset_
@@ -1094,20 +1094,37 @@
 ntities=
-True
+self.grouped_entities
 %0A       
@@ -2052,24 +2052,62 @@
 output = %22%22%0A
+        if self.grouped_entities:%0A    
         self
@@ -2150,16 +2150,17 @@
 %22, %22 %22),
+ 
 self.IOB
@@ -2204,16 +2204,128 @@
 on_ner%5D%0A
+        else:%0A            self.sent_ner = %5B(i%5B'word'%5D.replace(%22%3C_%3E%22, %22 %22), i%5B'entity'%5D) for i in self.json_ner%5D%0A
         

Update wangchanberta.py