UniTok 3.0.13__tar.gz → 3.1.1__tar.gz

{UniTok-3.0.13 → UniTok-3.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: UniTok
-Version: 3.0.13
+Version: 3.1.1
 Summary: Unified Tokenizer
 Home-page: https://github.com/Jyonn/UnifiedTokenizer
 Author: Jyonn Liu

{UniTok-3.0.13 → UniTok-3.1.1}/UniTok/column.py

@@ -1,9 +1,11 @@
-from UniTok.tok import BaseTok
+from typing import Type, Union
+
 from tqdm import tqdm
 
 from .global_setting import Global
 from .analysis.lengths import Lengths
-from .tok import IdTok
+from .tok import IdTok, BaseTok
+from .vocab import Vocab
 
 
 class SeqOperator:
@@ -42,11 +44,16 @@ class Column:
         tok (BaseTok): The tokenizer of the column.
         operator (SeqOperator): The operator of the column.
     """
-    def __init__(self, tok: BaseTok, name=None, operator: SeqOperator = None, **kwargs):
+    def __init__(self, tok: Union[BaseTok, Type[BaseTok]], name=None, operator: SeqOperator = None, **kwargs):
         self.tok = tok
         self.name = name or tok.vocab.name
         self.operator = operator
 
+        if isinstance(tok, type):
+            assert issubclass(tok, BaseTok)
+            assert name is not None, 'name must be set when tok is a class'
+            self.tok = tok(vocab=Vocab(name=name))
+
         if kwargs:
             if operator:
                 raise ValueError('operator and kwargs cannot be set at the same time')

{UniTok-3.0.13 → UniTok-3.1.1}/UniTok/meta.py

@@ -1,13 +1,13 @@
 import json
 import os
 import warnings
-from typing import List
+from typing import List, Union
 
 
 class Col:
     def __init__(self, name, voc=None, max_length=None, padding=None, vocab=None):
         self.name: str = name
-        self.voc: Voc = voc or vocab
+        self.voc: Union[Voc, str] = voc or vocab
         self.max_length = max_length
         self.padding = padding
         self.list = max_length is not None
@@ -29,7 +29,7 @@ class Voc:
     def __init__(self, name, size, cols, store_dir):
         self.name: str = name
         self.size: int = size
-        self.cols: List[Col] = cols
+        self.cols: List[Union[Col, str]] = cols
         self.store_dir = store_dir
 
     def __eq__(self, other):
@@ -46,7 +46,7 @@ class Voc:
         vocab = Vocab(name=self.name).load(self.store_dir)
         vocab.save(store_dir)
 
-    def merge(self, other):
+    def merge(self, other: 'Voc'):
         cols = self.cols.copy()
         for col in other.cols:
             for _col in cols:
@@ -71,13 +71,13 @@ class Meta:
 
         data = self.load()
         self.version = data['version']
-        self.cols = data.get('cols') or data['col_info']
-        self.vocs = data.get('vocs') or data['vocab_info']
+        cols = data.get('cols') or data['col_info']
+        vocs = data.get('vocs') or data['vocab_info']
         self.id_col = data['id_col']
 
         # build col-voc graph
-        self.cols = {col: Col(**self.cols[col], name=col) for col in self.cols}
-        self.vocs = {voc: Voc(**self.vocs[voc], name=voc, store_dir=self.store_dir) for voc in self.vocs}
+        self.cols = {col: Col(**cols[col], name=col) for col in cols}  # type: dict[str, Col]
+        self.vocs = {voc: Voc(**vocs[voc], name=voc, store_dir=self.store_dir) for voc in vocs}  # type: dict[str, Voc]
 
         # connect class objects
         for col in self.cols.values():

{UniTok-3.0.13 → UniTok-3.1.1}/UniTok/unidep.py

@@ -46,7 +46,7 @@ class UniDep:
             self.print('resize sample_size to', self._sample_size)
             self.sample_size = self._sample_size
 
-        self.vocabs = Vocabs()
+        self.vocabs = Vocabs()  # type: Union[Dict[str, Vocab], Vocabs]
         for vocab_name in self.vocs:
             self.vocabs.append(Vocab(name=vocab_name).load(self.store_dir))
         self.id2index = self.vocabs[self.id_voc.name].o2i
@@ -100,6 +100,11 @@ class UniDep:
         index = self._indexes[index]
         return self.pack_sample(index)
 
+    def __iter__(self):
+        """vocab obj list iterator"""
+        for i in range(len(self)):
+            yield self[i]
+
     def __len__(self):
         return self.sample_size
 

{UniTok-3.0.13 → UniTok-3.1.1}/UniTok/unitok.py

@@ -1,16 +1,14 @@
 import json
 import os
 import warnings
-from typing import Optional
+from typing import Optional, Type, Dict, Union
 
 import numpy as np
 import pandas as pd
 
 from .cols import Cols
 from .column import Column, IndexColumn
-from .tok.bert_tok import BertTok
-from .tok.ent_tok import EntTok
-from .tok.id_tok import IdTok
+from .tok import BaseTok, BertTok, EntTok, IdTok
 from .vocab import Vocab
 from .vocabs import Vocabs
 
@@ -60,7 +58,7 @@ class UniTok:
 
     def __init__(self):
         self.cols = Cols()
-        self.vocabs = Vocabs()
+        self.vocabs = Vocabs()  # type: Union[Dict[str, Vocab], Vocabs]
         self.id_col = None  # type: Optional[Column]
         self.data = None
 
@@ -70,10 +68,15 @@ class UniTok:
                       'use vocabs instead (will be removed in 4.x version)', DeprecationWarning)
         return self.vocabs
 
-    def add_col(self, col: Column):
+    def add_col(self, col: Union[Column, str], tok: Union[BaseTok, Type[BaseTok]] = None):
         """
         Declare a column in the DataFrame to be tokenized.
         """
+
+        if isinstance(col, str):
+            assert tok is not None, 'tok must be specified when col is a string'
+            col = Column(tok, name=col)
+
         if isinstance(col.tok, IdTok):
             if self.id_col:
                 raise ValueError(f'already exists id column {self.id_col.name} before adding {col.name}')

{UniTok-3.0.13 → UniTok-3.1.1}/UniTok/vocab.py

@@ -123,9 +123,15 @@ class Vocab:
         return True
 
     def __iter__(self):
+        """vocab obj list iterator"""
         for i in range(len(self)):
             yield self.i2o[i]
 
+    def __getitem__(self, item):
+        if isinstance(item, int):
+            return self.i2o[item]
+        return self.o2i[item]
+
     """
     Editable Methods
     """

{UniTok-3.0.13 → UniTok-3.1.1}/UniTok.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: UniTok
-Version: 3.0.13
+Version: 3.1.1
 Summary: Unified Tokenizer
 Home-page: https://github.com/Jyonn/UnifiedTokenizer
 Author: Jyonn Liu

{UniTok-3.0.13 → UniTok-3.1.1}/setup.py

@@ -6,7 +6,7 @@ long_description = (this_directory / "README.md").read_text()
 
 setup(
     name='UniTok',
-    version='3.0.13',
+    version='3.1.1',
     keywords=['token', 'tokenizer', 'bert'],
     description='Unified Tokenizer',
     long_description=long_description,