SinaTools 0.1.24__py2.py3-none-any.whl → 0.1.26__py2.py3-none-any.whl

{SinaTools-0.1.24.dist-info → SinaTools-0.1.26.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: SinaTools
-Version: 0.1.24
+Version: 0.1.26
 Summary: Open-source Python toolkit for Arabic Natural Understanding, allowing people to integrate it in their system workflow.
 Home-page: https://github.com/SinaLab/sinatools
 License: MIT license

{SinaTools-0.1.24.dist-info → SinaTools-0.1.26.dist-info}/RECORD

@@ -1,14 +1,14 @@
-SinaTools-0.1.24.data/data/sinatools/environment.yml,sha256=OzilhLjZbo_3nU93EQNUFX-6G5O3newiSWrwxvMH2Os,7231
-sinatools/VERSION,sha256=S6iCAVLWhyRA7MIGZk5mjvtI6v6w1_bWDqhs5ui2fDk,6
+SinaTools-0.1.26.data/data/sinatools/environment.yml,sha256=OzilhLjZbo_3nU93EQNUFX-6G5O3newiSWrwxvMH2Os,7231
+sinatools/VERSION,sha256=5E6i4X07Go6cKsVD3uEZkX9jXfyE05s7HlzVXSisTX8,6
 sinatools/__init__.py,sha256=bEosTU1o-FSpyytS6iVP_82BXHF2yHnzpJxPLYRbeII,135
 sinatools/environment.yml,sha256=OzilhLjZbo_3nU93EQNUFX-6G5O3newiSWrwxvMH2Os,7231
 sinatools/install_env.py,sha256=EODeeE0ZzfM_rz33_JSIruX03Nc4ghyVOM5BHVhsZaQ,404
 sinatools/sinatools.py,sha256=vR5AaF0iel21LvsdcqwheoBz0SIj9K9I_Ub8M8oA98Y,20
-sinatools/CLI/DataDownload/download_files.py,sha256=KG9W-Y5kJG_9yLUyo-cA33B5uO3avdZ5sSYUeW3wM6s,1960
+sinatools/CLI/DataDownload/download_files.py,sha256=VunXU_vAweKs7aS0FNM84N_2lhYT5T94Y8B3NWmGksg,2630
 sinatools/CLI/morphology/ALMA_multi_word.py,sha256=ZImJ1vtcpSHydI1BjJmK3KcMJbGBZX16kO4L6rxvBvA,2086
 sinatools/CLI/morphology/morph_analyzer.py,sha256=ieIM47QK9Nct3MtCS9uq3h2rZN5r4qNhsLmlVeE6wiE,3503
-sinatools/CLI/ner/corpus_entity_extractor.py,sha256=_o0frMSgpsFVXPoztS3mQTK7LjHsgzUv9gfs6iJL424,4024
-sinatools/CLI/ner/entity_extractor.py,sha256=zn0Jd37BEDE1wHE5HOAK0_N2tURAznFNj7WDd6WGLIw,2932
+sinatools/CLI/ner/corpus_entity_extractor.py,sha256=Da-DHFrqT6if7w6WnodB4TBE5ze3DJYjb2Mmju_Qd7g,4034
+sinatools/CLI/ner/entity_extractor.py,sha256=IiTioe0px0aJ1E58FrDVa2yNgM8Ie4uS2LZKK_z2Qn4,2942
 sinatools/CLI/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 sinatools/CLI/utils/arStrip.py,sha256=NLyp8vOu2xv80tL9jiKRvyptmbkRZVg-wcAr-9YyvNY,3264
 sinatools/CLI/utils/corpus_tokenizer.py,sha256=nH0T4h6urr_0Qy6-wN3PquOtnwybj0REde5Ts_OE4U8,1650
@@ -20,7 +20,7 @@ sinatools/CLI/utils/sentence_tokenizer.py,sha256=Wli8eiDbWSd_Z8UKpu_JkaS8jImowa1
 sinatools/CLI/utils/text_dublication_detector.py,sha256=dW70O5O20GxeUDDF6zVYn52wWLmJF-HBZgvqIeVL2rQ,1661
 sinatools/CLI/utils/text_transliteration.py,sha256=vz-3kxWf8pNYVCqNAtBAiA6u_efrS5NtWT-ofN1NX6I,2014
 sinatools/DataDownload/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-sinatools/DataDownload/downloader.py,sha256=sLmVvnr3mG_tqvGCggzxwsi2sixlKlgCbMnZhCclSpg,6390
+sinatools/DataDownload/downloader.py,sha256=F-SV-0mbYMYFSNCx8FoAYXhn0X1j0dF37PTLU0nUBVg,6482
 sinatools/arabert/__init__.py,sha256=ely2PttjgSv7vKdzskuD1rtK_l_UOpmxJSz8isrveD0,16
 sinatools/arabert/preprocess.py,sha256=qI0FsuMTOzdRlYGCtLrjpXgikNElUZPv9bnjaKDZKJ4,33024
 sinatools/arabert/arabert/__init__.py,sha256=KbSAH-XqbRygn0y59m5-ZYOLXgpT1gSgE3F-qd4rKEc,627
@@ -80,7 +80,7 @@ sinatools/ner/__init__.py,sha256=gSs0x6veWJ8j3_iOs79tynBd_hJP0t44CGpJ0xzoiW4,104
 sinatools/ner/data.py,sha256=lvOW86dXse8SC75Q0supQaE0rrRffoxNjIA0Qbv5WZY,4354
 sinatools/ner/data_format.py,sha256=7Yt0aOicOn9_YuuyCkM_IYi_rgjGYxR9bCuUaNGM73o,4341
 sinatools/ner/datasets.py,sha256=mG1iwqSm3lXCFHLqE-b4wNi176cpuzNBz8tKaBU6z6M,5059
-sinatools/ner/entity_extractor.py,sha256=k0Yvvg_aknINkFSdqOgG1KulS0UIo-W0qycv9J2MtNo,2273
+sinatools/ner/entity_extractor.py,sha256=yQnfayT03qAnQ4FBdBFhvl8M2pgIttrdWSWE9wgO2LI,1876
 sinatools/ner/helpers.py,sha256=dnOoDY5JMyOLTUWVIZLMt8mBn2IbWlVaqHhQyjs1voo,2343
 sinatools/ner/metrics.py,sha256=Irz6SsIvpOzGIA2lWxrEV86xnTnm0TzKm9SUVT4SXUU,2734
 sinatools/ner/transforms.py,sha256=vti3mDdi-IRP8i0aTQ37QqpPlP9hdMmJ6_bAMa0uL-s,4871
@@ -91,7 +91,7 @@ sinatools/ner/nn/BaseModel.py,sha256=3GmujQasTZZunOBuFXpY2p1W8W256iI_Uu4hxhOY2Z0
 sinatools/ner/nn/BertNestedTagger.py,sha256=_fwAn1kiKmXe6m5y16Ipty3kvXIEFEmiUq74Ad1818U,1219
 sinatools/ner/nn/BertSeqTagger.py,sha256=dFcBBiMw2QCWsyy7aQDe_PS3aRuNn4DOxKIHgTblFvc,504
 sinatools/ner/nn/__init__.py,sha256=UgQD_XLNzQGBNSYc_Bw1aRJZjq4PJsnMT1iZwnJemqE,170
-sinatools/ner/trainers/BaseTrainer.py,sha256=oZgFJW-CawfCKT5gtaBHA7Q7XjNfiyqM62KnFsgVzPU,3919
+sinatools/ner/trainers/BaseTrainer.py,sha256=Ifz4SeTxJwVn1_uWZ3I9KbcSo2hLPN3ojsIYuoKE9wE,4050
 sinatools/ner/trainers/BertNestedTrainer.py,sha256=Pb4O2WeBmTvV3hHMT6DXjxrTzgtuh3OrKQZnogYy8RQ,8429
 sinatools/ner/trainers/BertTrainer.py,sha256=B_uVtUwfv_eFwMMPsKQvZgW_ZNLy6XEsX5ePR0s8d-k,6433
 sinatools/ner/trainers/__init__.py,sha256=UDok8pDDpYOpwRBBKVLKaOgSUlmqqb-zHZI1p0xPxzI,188
@@ -110,13 +110,13 @@ sinatools/utils/text_transliteration.py,sha256=NQoXrxI-h0UXnvVtDA3skNJduxIy0IW26
 sinatools/utils/tokenizer.py,sha256=QHyrVqJA_On4rKxexiWR2ovq4pI1-u6iZkdhRbK9tew,6676
 sinatools/utils/tokenizers_words.py,sha256=efNfOil9qDNVJ9yynk_8sqf65PsL-xtsHG7y2SZCkjQ,656
 sinatools/wsd/__init__.py,sha256=yV-SQSCzSrjbNkciMbDCqzGZ_EESchL7rlJk56uibVI,309
-sinatools/wsd/disambiguator.py,sha256=8HrVAGpEQyrzwiuEreLX9X82WSL-U2Aeca0ttrtIw2Y,19998
+sinatools/wsd/disambiguator.py,sha256=43Iq7NTZsiYWGFg-NUDrQuJKO1NT9QOnfBPB10IOJNs,19828
 sinatools/wsd/settings.py,sha256=6XflVTFKD8SVySX9Wj7zYQtV26WDTcQ2-uW8-gDNHKE,747
 sinatools/wsd/wsd.py,sha256=gHIBUFXegoY1z3rRnIlK6TduhYq2BTa_dHakOjOlT4k,4434
-SinaTools-0.1.24.dist-info/AUTHORS.rst,sha256=aTWeWlIdfLi56iLJfIUAwIrmqDcgxXKLji75_Fjzjyg,174
-SinaTools-0.1.24.dist-info/LICENSE,sha256=uwsKYG4TayHXNANWdpfMN2lVW4dimxQjA_7vuCVhD70,1088
-SinaTools-0.1.24.dist-info/METADATA,sha256=TS_IfzeMqZsoClo4KPnnhsTHbuo8sWNBXB2ByHkrY_M,953
-SinaTools-0.1.24.dist-info/WHEEL,sha256=6T3TYZE4YFi2HTS1BeZHNXAi8N52OZT4O-dJ6-ome_4,116
-SinaTools-0.1.24.dist-info/entry_points.txt,sha256=ZwZLolnWog2fjdDrfaHNHob8SE_YtMbD6ayzsOzItxs,1234
-SinaTools-0.1.24.dist-info/top_level.txt,sha256=8tNdPTeJKw3TQCaua8IJIx6N6WpgZZmVekf1OdBNJpE,10
-SinaTools-0.1.24.dist-info/RECORD,,
+SinaTools-0.1.26.dist-info/AUTHORS.rst,sha256=aTWeWlIdfLi56iLJfIUAwIrmqDcgxXKLji75_Fjzjyg,174
+SinaTools-0.1.26.dist-info/LICENSE,sha256=uwsKYG4TayHXNANWdpfMN2lVW4dimxQjA_7vuCVhD70,1088
+SinaTools-0.1.26.dist-info/METADATA,sha256=jqsARSXI1Z0hT9-ev6ewzZeNH_H350lv_c2oav_SKWg,953
+SinaTools-0.1.26.dist-info/WHEEL,sha256=6T3TYZE4YFi2HTS1BeZHNXAi8N52OZT4O-dJ6-ome_4,116
+SinaTools-0.1.26.dist-info/entry_points.txt,sha256=ZwZLolnWog2fjdDrfaHNHob8SE_YtMbD6ayzsOzItxs,1234
+SinaTools-0.1.26.dist-info/top_level.txt,sha256=8tNdPTeJKw3TQCaua8IJIx6N6WpgZZmVekf1OdBNJpE,10
+SinaTools-0.1.26.dist-info/RECORD,,

sinatools/CLI/DataDownload/download_files.py

@@ -40,7 +40,7 @@ from sinatools.DataDownload.downloader import urls
 
 def main():
     parser = argparse.ArgumentParser(description="Download files from specified URLs.")
-    parser.add_argument('-f', '--files', nargs="*", choices=urls.keys(),
+    parser.add_argument('-f', '--files', nargs="*",
                         help="Names of the files to download. Available files are: "
                              f"{', '.join(urls.keys())}. If no file is specified, all files will be downloaded.")
     
@@ -50,8 +50,23 @@ def main():
 
     if args.files:
         for file in args.files:
-            url = urls[file]
-            download_file(url)
+            print("file: ", file)
+            if file == "wsd":
+                download_file(urls["morph"])
+                download_file(urls["ner"])
+                download_file(urls["wsd_model"])
+                download_file(urls["wsd_tokenizer"])
+                download_file(urls["glosses_dic"])
+                download_file(urls["five_grams"])
+                download_file(urls["four_grams"])
+                download_file(urls["three_grams"])
+                download_file(urls["two_grams"])
+            elif file == "synonyms":
+                download_file(urls["synonyms_level2"])
+                download_file(urls["synonyms_level3"])
+            else:
+               url = urls[file]
+               download_file(url)
     else:
         download_files()
 

sinatools/CLI/ner/corpus_entity_extractor.py

@@ -20,7 +20,7 @@ def jsons_to_list_of_lists(json_list):
     return [[d['token'], d['tags']] for d in json_list]
 
 def combine_tags(sentence):
-    output = jsons_to_list_of_lists(extract(sentence))
+    output = jsons_to_list_of_lists(extract(sentence, "nested"))
     return [word[1] for word in output]
 
 

sinatools/CLI/ner/entity_extractor.py

@@ -46,7 +46,7 @@ def jsons_to_list_of_lists(json_list):
     return [[d['token'], d['tags']] for d in json_list]
 
 def combine_tags(sentence):
-    output = jsons_to_list_of_lists(extract(sentence))
+    output = jsons_to_list_of_lists(extract(sentence, "nested"))
     return [word[1] for word in output]
 
 

sinatools/DataDownload/downloader.py

@@ -95,37 +95,41 @@ def download_file(url, dest_path=get_appdatadir()):
     print(filename)
     headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3'}
 
-    # try:
-    with requests.get(url, headers=headers, stream=True) as r:
-        r.raise_for_status()
-        with open(file_path, 'wb') as f:
-            total_size = int(r.headers.get('content-length', 0))
-            block_size = 8192
-            progress_bar = tqdm(total=total_size, unit='iB', unit_scale=True)
-            for chunk in r.iter_content(chunk_size=block_size):
-                if chunk:
-                    f.write(chunk)
-                    progress_bar.update(len(chunk))
-            progress_bar.close()
-    # Check the file type and extract accordingly
-    file_extension = os.path.splitext(file_path)[1]
-    extracted_folder_name = os.path.splitext(file_path)[0]
-    
-    if file_extension == '.zip':
-        extract_zip(file_path, extracted_folder_name)
-    elif file_extension == '.gz':
-        extract_tar(file_path, extracted_folder_name)
-    elif file_extension =='.pickle':
-        print(f'Done: {file_extension}')
-    else:
-        print(f'Unsupported file type for extraction: {file_extension}')
-    return file_path
-
-    # except requests.exceptions.HTTPError as e:
-        # if e.response.status_code == 403:
-            # print(f'Error 403: Forbidden. The requested file URL {url} could not be downloaded due to insufficient permissions. Please check the URL and try again.')
-        # else:
-            # print('An error occurred while downloading the file:', e)
+    try:
+        with requests.get(url, headers=headers, stream=True) as r:
+            r.raise_for_status()
+            with open(file_path, 'wb') as f:
+                total_size = int(r.headers.get('content-length', 0))
+                block_size = 8192
+                progress_bar = tqdm(total=total_size, unit='iB', unit_scale=True)
+                for chunk in r.iter_content(chunk_size=block_size):
+                    if chunk:
+                        f.write(chunk)
+                        progress_bar.update(len(chunk))
+                progress_bar.close()
+
+        # Check the file type and extract accordingly
+        file_extension = os.path.splitext(file_path)[1]
+        extracted_folder_name = os.path.splitext(file_path)[0]
+        
+        if file_extension == '.zip':
+            extract_zip(file_path, extracted_folder_name)
+        elif file_extension == '.gz':
+
+            extract_tar(file_path, extracted_folder_name)
+        elif file_extension =='.pickle':
+            print(f'Done: {file_extension}')
+
+        else:
+            print(f'Unsupported file type for extraction: {file_extension}')
+
+        return file_path
+
+    except requests.exceptions.HTTPError as e:
+        if e.response.status_code == 403:
+            print(f'Error 403: Forbidden. The requested file URL {url} could not be downloaded due to insufficient permissions. Please check the URL and try again.')
+        else:
+            print('An error occurred while downloading the file:', e)
 
 def extract_zip(file_path, extracted_folder_name):
     """

sinatools/VERSION

@@ -1 +1 @@
-0.1.24
+0.1.26

sinatools/ner/entity_extractor.py

@@ -3,43 +3,31 @@ from collections import namedtuple
 from sinatools.ner.data_format import get_dataloaders, text2segments
 from . import tagger, tag_vocab, train_config
 
-def extract(text, batch_size=32):
-    """
-    This method processes an input text and returns named entites for each token within the text, based on the specified batch size. As follows:
 
-    Args:
-        text (:obj:`str`): The Arabic text to be tagged.
-        batch_size (int, optional): Batch size for inference. Default is 32.
-
-    Returns:
-        list (:obj:`list`): A list of JSON objects, where each JSON could be contains:
-        token: The token from the original text.
-        NER tag: The label pairs for each segment.
-
-    **Example:**
-
-     .. highlight:: python
-     .. code-block:: python
+def convert_nested_to_flat(nested_tags):
+    flat_tags = []
+    
+    for entry in nested_tags:
+        word = entry['token']
+        tags = entry['tags'].split()
+        
+        # Initialize with the first tag in the sequence
+        flat_tag = tags[0]
+        
+        for tag in tags[1:]:
+            # Check if the tag is an "I-" tag, indicating continuation of an entity
+            if tag.startswith('I-'):
+                flat_tag = tag
+                break
+        
+        flat_tags.append({
+            'token': word,
+            'tags': flat_tag
+        })
+    
+    return flat_tags
 
-        from sinatools.ner.entity_extractor import extract
-        extract('ذهب محمد إلى جامعة بيرزيت')
-        [{
-            "word":"ذهب",
-            "tags":"O"
-          },{
-            "word":"محمد",
-            "tags":"B-PERS"
-          },{
-            "word":"إلى",
-            "tags":"O"
-          },{
-            "word":"جامعة",
-            "tags":"B-ORG"
-          },{
-            "word":"بيرزيت",
-            "tags":"B-GPE I-ORG"
-        }]
-    """    
+def extract(text, ner_method):
     
     dataset, token_vocab = text2segments(text)
 
@@ -50,7 +38,7 @@ def extract(text, batch_size=32):
         (dataset,),
         vocab,
         train_config.data_config,
-        batch_size=batch_size,
+        batch_size=32,
         shuffle=(False,),
     )[0]
 
@@ -69,4 +57,7 @@ def extract(text, batch_size=32):
             else:
                segments_list["tags"] = ' '.join(list_of_tags)
             segments_lists.append(segments_list)  
+    
+    if ner_method == "flat":
+      segments_lists = convert_nested_to_flat(segments_lists)          
     return segments_lists

sinatools/ner/trainers/BaseTrainer.py

@@ -1,117 +1,117 @@
-import os
-import torch
-import logging
-import natsort
-import glob
-
-logger = logging.getLogger(__name__)
-
-
-class BaseTrainer:
-    def __init__(
-        self,
-        model=None,
-        max_epochs=50,
-        optimizer=None,
-        scheduler=None,
-        loss=None,
-        train_dataloader=None,
-        val_dataloader=None,
-        test_dataloader=None,
-        log_interval=10,
-        summary_writer=None,
-        output_path=None,
-        clip=5,
-        patience=5
-    ):
-        self.model = model
-        self.max_epochs = max_epochs
-        self.train_dataloader = train_dataloader
-        self.val_dataloader = val_dataloader
-        self.test_dataloader = test_dataloader
-        self.optimizer = optimizer
-        self.scheduler = scheduler
-        self.loss = loss
-        self.log_interval = log_interval
-        self.summary_writer = summary_writer
-        self.output_path = output_path
-        self.current_timestep = 0
-        self.current_epoch = 0
-        self.clip = clip
-        self.patience = patience
-
-    def tag(self, dataloader, is_train=True):
-        """
-        Given a dataloader containing segments, predict the tags
-        :param dataloader: torch.utils.data.DataLoader
-        :param is_train: boolean - True for training model, False for evaluation
-        :return: Iterator
-                    subwords (B x T x NUM_LABELS)- torch.Tensor - BERT subword ID
-                    gold_tags (B x T x NUM_LABELS) - torch.Tensor - ground truth tags IDs
-                    tokens - List[arabiner.data.dataset.Token] - list of tokens
-                    valid_len (B x 1) - int - valiud length of each sequence
-                    logits (B x T x NUM_LABELS) - logits for each token and each tag
-        """
-        for subwords, gold_tags, tokens, valid_len in dataloader:
-            self.model.train(is_train)
-
-            if torch.cuda.is_available():
-                subwords = subwords.cuda()
-                gold_tags = gold_tags.cuda()
-
-            if is_train:
-                self.optimizer.zero_grad()
-                logits = self.model(subwords)
-            else:
-                with torch.no_grad():
-                    logits = self.model(subwords)
-
-            yield subwords, gold_tags, tokens, valid_len, logits
-
-    def segments_to_file(self, segments, filename):
-        """
-        Write segments to file
-        :param segments: [List[arabiner.data.dataset.Token]] - list of list of tokens
-        :param filename: str - output filename
-        :return: None
-        """
-        with open(filename, "w") as fh:
-            results = "\n\n".join(["\n".join([t.__str__() for t in segment]) for segment in segments])
-            fh.write("Token\tGold Tag\tPredicted Tag\n")
-            fh.write(results)
-            logging.info("Predictions written to %s", filename)
-
-    def save(self):
-        """
-        Save model checkpoint
-        :return:
-        """
-        filename = os.path.join(
-            self.output_path,
-            "checkpoints",
-            "checkpoint_{}.pt".format(self.current_epoch),
-        )
-
-        checkpoint = {
-            "model": self.model.state_dict(),
-            "optimizer": self.optimizer.state_dict(),
-            "epoch": self.current_epoch
-        }
-
-        logger.info("Saving checkpoint to %s", filename)
-        torch.save(checkpoint, filename)
-
-    def load(self, checkpoint_path):
-        """
-        Load model checkpoint
-        :param checkpoint_path: str - path/to/checkpoints
-        :return: None
-        """
-        checkpoint_path = natsort.natsorted(glob.glob(f"{checkpoint_path}/checkpoint_*.pt"))
-        checkpoint_path = checkpoint_path[-1]
-
-        logger.info("Loading checkpoint %s", checkpoint_path)
-
-        device = None if torch.cuda.is_available() else torch.device('cpu')
-        checkpoint = torch.load(checkpoint_path, map_location=device)
-        self.model.load_state_dict(checkpoint["model"])
+import os
+import torch
+import logging
+import natsort
+import glob
+
+logger = logging.getLogger(__name__)
+
+
+class BaseTrainer:
+    def __init__(
+        self,
+        model=None,
+        max_epochs=50,
+        optimizer=None,
+        scheduler=None,
+        loss=None,
+        train_dataloader=None,
+        val_dataloader=None,
+        test_dataloader=None,
+        log_interval=10,
+        summary_writer=None,
+        output_path=None,
+        clip=5,
+        patience=5
+    ):
+        self.model = model
+        self.max_epochs = max_epochs
+        self.train_dataloader = train_dataloader
+        self.val_dataloader = val_dataloader
+        self.test_dataloader = test_dataloader
+        self.optimizer = optimizer
+        self.scheduler = scheduler
+        self.loss = loss
+        self.log_interval = log_interval
+        self.summary_writer = summary_writer
+        self.output_path = output_path
+        self.current_timestep = 0
+        self.current_epoch = 0
+        self.clip = clip
+        self.patience = patience
+
+    def tag(self, dataloader, is_train=True):
+        """
+        Given a dataloader containing segments, predict the tags
+        :param dataloader: torch.utils.data.DataLoader
+        :param is_train: boolean - True for training model, False for evaluation
+        :return: Iterator
+                    subwords (B x T x NUM_LABELS)- torch.Tensor - BERT subword ID
+                    gold_tags (B x T x NUM_LABELS) - torch.Tensor - ground truth tags IDs
+                    tokens - List[arabiner.data.dataset.Token] - list of tokens
+                    valid_len (B x 1) - int - valiud length of each sequence
+                    logits (B x T x NUM_LABELS) - logits for each token and each tag
+        """
+        for subwords, gold_tags, tokens, valid_len in dataloader:
+            self.model.train(is_train)
+
+            if torch.cuda.is_available():
+                subwords = subwords.cuda()
+                gold_tags = gold_tags.cuda()
+
+            if is_train:
+                self.optimizer.zero_grad()
+                logits = self.model(subwords)
+            else:
+                with torch.no_grad():
+                    logits = self.model(subwords)
+
+            yield subwords, gold_tags, tokens, valid_len, logits
+
+    def segments_to_file(self, segments, filename):
+        """
+        Write segments to file
+        :param segments: [List[arabiner.data.dataset.Token]] - list of list of tokens
+        :param filename: str - output filename
+        :return: None
+        """
+        with open(filename, "w") as fh:
+            results = "\n\n".join(["\n".join([t.__str__() for t in segment]) for segment in segments])
+            fh.write("Token\tGold Tag\tPredicted Tag\n")
+            fh.write(results)
+            logging.info("Predictions written to %s", filename)
+
+    def save(self):
+        """
+        Save model checkpoint
+        :return:
+        """
+        filename = os.path.join(
+            self.output_path,
+            "checkpoints",
+            "checkpoint_{}.pt".format(self.current_epoch),
+        )
+
+        checkpoint = {
+            "model": self.model.state_dict(),
+            "optimizer": self.optimizer.state_dict(),
+            "epoch": self.current_epoch
+        }
+
+        logger.info("Saving checkpoint to %s", filename)
+        torch.save(checkpoint, filename)
+
+    def load(self, checkpoint_path):
+        """
+        Load model checkpoint
+        :param checkpoint_path: str - path/to/checkpoints
+        :return: None
+        """
+        checkpoint_path = natsort.natsorted(glob.glob(f"{checkpoint_path}/checkpoint_*.pt"))
+        checkpoint_path = checkpoint_path[-1]
+
+        logger.info("Loading checkpoint %s", checkpoint_path)
+
+        device = None if torch.cuda.is_available() else torch.device('cpu')
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+        self.model.load_state_dict(checkpoint["model"], strict=False)

sinatools/wsd/disambiguator.py

@@ -217,7 +217,7 @@ def jsons_to_list_of_lists(json_list):
 def find_named_entities(string):
    found_entities = []
    
-   ner_entites = extract(string)
+   ner_entites = extract(string, "nested")
    list_of_entites = jsons_to_list_of_lists(ner_entites)
    entites = distill_entities(list_of_entites)
    
@@ -288,17 +288,17 @@ def disambiguate_glosses_using_SALMA(glosses, Diac_lemma, Undiac_lemma, word, se
       concept_id, gloss = GlossPredictor(Diac_lemma, Undiac_lemma,word,sentence,glosses_dictionary)
 
       my_json = {}    
-      my_json['Concept_id'] = concept_id
+      my_json['concept_id'] = concept_id
     #   my_json['Gloss'] = gloss
       my_json['word'] = word
-      my_json['Undiac_lemma'] = Undiac_lemma
-      my_json['Diac_lemma'] = Diac_lemma
+      #my_json['Undiac_lemma'] = Undiac_lemma
+      my_json['lemma'] = Diac_lemma
       return my_json
    else:
       my_json = {}    
       my_json['word'] = word
-      my_json['Undiac_lemma'] = Undiac_lemma
-      my_json['Diac_lemma'] = Diac_lemma
+      #my_json['Undiac_lemma'] = Undiac_lemma
+      my_json['lemma'] = Diac_lemma
       return my_json
 
 
@@ -405,26 +405,26 @@ def disambiguate_glosses_main(word, sentence):
    if concept_count == 0:
       my_json = {}    
       my_json['word'] = word['word']
-      my_json['Diac_lemma'] = word['Diac_lemma']
-      my_json['Undiac_lemma'] = word['Undiac_lemma']
+      my_json['lemma'] = word['Diac_lemma']
+      #my_json['Undiac_lemma'] = word['Undiac_lemma']
       return my_json
    elif concept_count == 1:
       my_json = {}    
       my_json['word'] = word['word']
       glosses = word['glosses'][0]
     #   my_json['Gloss'] = glosses['gloss']
-      my_json['Concept_id'] = glosses['concept_id']
-      my_json['Diac_lemma'] = word['Diac_lemma']
-      my_json['Undiac_lemma'] = word['Undiac_lemma']
+      my_json['concept_id'] = glosses['concept_id']
+      my_json['lemma'] = word['Diac_lemma']
+      #my_json['Undiac_lemma'] = word['Undiac_lemma']
       return my_json
    elif concept_count == '*':
       my_json = {}    
       my_json['word'] = word['word']
       glosses = word['glosses'][0]
       my_json['Gloss'] = glosses['gloss']
-      my_json['Concept_id'] = glosses['concept_id']
-      my_json['Diac_lemma'] = word['Diac_lemma']
-      my_json['Undiac_lemma'] = word['Undiac_lemma']
+      my_json['concept_id'] = glosses['concept_id']
+      my_json['lemma'] = word['Diac_lemma']
+      #my_json['Undiac_lemma'] = word['Undiac_lemma']
       return my_json
    else:   
       input_word = word['word']
@@ -477,21 +477,18 @@ def disambiguate(sentence):
         #output
          [
              {
-                 "Concept_id": "303019218",
+                 "concept_id": "303019218",
                  "word": "ذهبت",
-                 "Undiac_lemma": "ذهب",
-                 "Diac_lemma": "ذَهَبَ۪ 1"
+                 "lemma": "ذَهَبَ۪ 1"
              },
              {
                  "word": "إلى",
-                 "Diac_lemma": إِلَى 1,
-                 "Undiac_lemma": "الى"
+                 "lemma": "إِلَى 1"
              },
              {
                  "word": "جامعة بيرزيت",
-                 "Concept_id": "334000099",
-                 "Diac_lemma": جامِعَة بيرزَيت,
-                 "Undiac_lemma": "جامعة بيرزيت"
+                 "concept_id": "334000099",
+                 "lemma": "جامِعَة بيرزَيت"
              }
          ]
     """