safetensors update (#3)

- safetensors update (50282f3d0697ea21779c53c244314517c6f80f78)
- safetensors model added (04ff336e3e47280791bd8e300468537e6179b03e)
- readme update (044cecb68768b5b450f58c726f4c2403838bd688)
- some formatting (a614a1af09ec9700164d5b56a9dd9f187360e2a4)
- inference false (bb4e784327033e0e11d5d61184d258a46bf93c0c)

.gitattributes

@@ -29,4 +29,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zstandard filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
-pytorch_model.bin filter=lfs diff=lfs merge=lfs -text
+model.safetensors filter=lfs diff=lfs merge=lfs -text

README.md

@@ -5,8 +5,8 @@ tags:
 - PyTorch
 - Transformers
 license: apache-2.0
-widget:
-- text: "sbert punc case расставляет точки запятые и знаки вопроса вам нравится"
+base_model: sberbank-ai/sbert_large_nlu_ru
+inference: false
 ---
 
 # SbertPuncCase

config.json

@@ -55,8 +55,8 @@
   "pooler_type": "first_token_transform",
   "position_embedding_type": "absolute",
   "torch_dtype": "float16",
-  "transformers_version": "4.20.1",
+  "transformers_version": "4.36.2",
   "type_vocab_size": 2,
   "use_cache": true,
   "vocab_size": 120138
-}
+}

pytorch_model.bin → model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:a0928d162fd53b902c8aa1704cb29f904d777398c152e0c6a5cdc676d6cf397c
-size 851804225
+oid sha256:06173fe7aed01a3a58385f6e724502f634909bee0818f6a8514b9eb6eb869be8
+size 851791402

sbert_punc_case_ru/sbertpunccase.py

@@ -8,62 +8,66 @@ import numpy as np
 from transformers import AutoTokenizer, AutoModelForTokenClassification
 
 # Прогнозируемые знаки препинания
-PUNK_MAPPING = {'.': 'PERIOD', ',': 'COMMA', '?': 'QUESTION'}
+PUNK_MAPPING = {".": "PERIOD", ",": "COMMA", "?": "QUESTION"}
 
 # Прогнозируемый регистр LOWER - нижний регистр, UPPER - верхний регистр для первого символа,
 # UPPER_TOTAL - верхний регистр для всех символов
-LABELS_CASE = ['LOWER', 'UPPER', 'UPPER_TOTAL']
+LABELS_CASE = ["LOWER", "UPPER", "UPPER_TOTAL"]
 # Добавим в пунктуацию метку O означающий отсутсвие пунктуации
-LABELS_PUNC = ['O'] + list(PUNK_MAPPING.values())
+LABELS_PUNC = ["O"] + list(PUNK_MAPPING.values())
 
 # Сформируем метки на основе комбинаций регистра и пунктуации
 LABELS_list = []
 for case in LABELS_CASE:
     for punc in LABELS_PUNC:
-        LABELS_list.append(f'{case}_{punc}')
-LABELS = {label: i+1 for i, label in enumerate(LABELS_list)}
-LABELS['O'] = -100
+        LABELS_list.append(f"{case}_{punc}")
+LABELS = {label: i + 1 for i, label in enumerate(LABELS_list)}
+LABELS["O"] = -100
 INVERSE_LABELS = {i: label for label, i in LABELS.items()}
 
-LABEL_TO_PUNC_LABEL = {label: label.split('_')[-1] for label in LABELS.keys() if label != 'O'}
-LABEL_TO_CASE_LABEL = {label: '_'.join(label.split('_')[:-1]) for label in LABELS.keys() if label != 'O'}
+LABEL_TO_PUNC_LABEL = {
+    label: label.split("_")[-1] for label in LABELS.keys() if label != "O"
+}
+LABEL_TO_CASE_LABEL = {
+    label: "_".join(label.split("_")[:-1]) for label in LABELS.keys() if label != "O"
+}
 
 
 def token_to_label(token, label):
     if type(label) == int:
         label = INVERSE_LABELS[label]
-    if label == 'LOWER_O':
+    if label == "LOWER_O":
         return token
-    if label == 'LOWER_PERIOD':
-        return token + '.'
-    if label == 'LOWER_COMMA':
-        return token + ','
-    if label == 'LOWER_QUESTION':
-        return token + '?'
-    if label == 'UPPER_O':
+    if label == "LOWER_PERIOD":
+        return token + "."
+    if label == "LOWER_COMMA":
+        return token + ","
+    if label == "LOWER_QUESTION":
+        return token + "?"
+    if label == "UPPER_O":
         return token.capitalize()
-    if label == 'UPPER_PERIOD':
-        return token.capitalize() + '.'
-    if label == 'UPPER_COMMA':
-        return token.capitalize() + ','
-    if label == 'UPPER_QUESTION':
-        return token.capitalize() + '?'
-    if label == 'UPPER_TOTAL_O':
+    if label == "UPPER_PERIOD":
+        return token.capitalize() + "."
+    if label == "UPPER_COMMA":
+        return token.capitalize() + ","
+    if label == "UPPER_QUESTION":
+        return token.capitalize() + "?"
+    if label == "UPPER_TOTAL_O":
         return token.upper()
-    if label == 'UPPER_TOTAL_PERIOD':
-        return token.upper() + '.'
-    if label == 'UPPER_TOTAL_COMMA':
-        return token.upper() + ','
-    if label == 'UPPER_TOTAL_QUESTION':
-        return token.upper() + '?'
-    if label == 'O':
+    if label == "UPPER_TOTAL_PERIOD":
+        return token.upper() + "."
+    if label == "UPPER_TOTAL_COMMA":
+        return token.upper() + ","
+    if label == "UPPER_TOTAL_QUESTION":
+        return token.upper() + "?"
+    if label == "O":
         return token
 
 
-def decode_label(label, classes='all'):
-    if classes == 'punc':
+def decode_label(label, classes="all"):
+    if classes == "punc":
         return LABEL_TO_PUNC_LABEL[INVERSE_LABELS[label]]
-    if classes == 'case':
+    if classes == "case":
         return LABEL_TO_CASE_LABEL[INVERSE_LABELS[label]]
     else:
         return INVERSE_LABELS[label]
@@ -76,14 +80,12 @@ class SbertPuncCase(nn.Module):
     def __init__(self):
         super().__init__()
 
-        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO,
-                                                       strip_accents=False)
+        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO, strip_accents=False)
         self.model = AutoModelForTokenClassification.from_pretrained(MODEL_REPO)
         self.model.eval()
 
     def forward(self, input_ids, attention_mask):
-        return self.model(input_ids=input_ids,
-                          attention_mask=attention_mask)
+        return self.model(input_ids=input_ids, attention_mask=attention_mask)
 
     def punctuate(self, text):
         text = text.strip().lower()
@@ -94,10 +96,23 @@ class SbertPuncCase(nn.Module):
         tokenizer_output = self.tokenizer(words, is_split_into_words=True)
 
         if len(tokenizer_output.input_ids) > 512:
-            return ' '.join([self.punctuate(' '.join(text_part)) for text_part in np.array_split(words, 2)])
-
-        predictions = self(torch.tensor([tokenizer_output.input_ids], device=self.model.device),
-                           torch.tensor([tokenizer_output.attention_mask], device=self.model.device)).logits.cpu().data.numpy()
+            return " ".join(
+                [
+                    self.punctuate(" ".join(text_part))
+                    for text_part in np.array_split(words, 2)
+                ]
+            )
+
+        predictions = (
+            self(
+                torch.tensor([tokenizer_output.input_ids], device=self.model.device),
+                torch.tensor(
+                    [tokenizer_output.attention_mask], device=self.model.device
+                ),
+            )
+            .logits.cpu()
+            .data.numpy()
+        )
         predictions = np.argmax(predictions, axis=2)
 
         # decode punctuation and casing
@@ -108,16 +123,31 @@ class SbertPuncCase(nn.Module):
             label_id = predictions[0][label_pos]
             label = decode_label(label_id)
             splitted_text.append(token_to_label(word, label))
-        capitalized_text = ' '.join(splitted_text)
+        capitalized_text = " ".join(splitted_text)
         return capitalized_text
 
 
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser("Punctuation and case restoration model sbert_punc_case_ru")
-    parser.add_argument("-i", "--input", type=str, help="text to restore", default='sbert punc case расставляет точки запятые и знаки вопроса вам нравится')
-    parser.add_argument("-d", "--device", type=str, help="run model on cpu or gpu", choices=['cpu', 'cuda'], default='cpu')
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        "Punctuation and case restoration model sbert_punc_case_ru"
+    )
+    parser.add_argument(
+        "-i",
+        "--input",
+        type=str,
+        help="text to restore",
+        default="sbert punc case расставляет точки запятые и знаки вопроса вам нравится",
+    )
+    parser.add_argument(
+        "-d",
+        "--device",
+        type=str,
+        help="run model on cpu or gpu",
+        choices=["cpu", "cuda"],
+        default="cpu",
+    )
     args = parser.parse_args()
     print(f"Source text:   {args.input}\n")
     sbertpunc = SbertPuncCase().to(args.device)
     punctuated_text = sbertpunc.punctuate(args.input)
-    print(f"Restored text: {punctuated_text}")
+    print(f"Restored text: {punctuated_text}")

setup.py

@@ -1,19 +1,24 @@
 from distutils.core import setup
 
-setup(name='sbert_punc_case_ru',
-      version='0.1',
-      description='Punctuation and Case Restoration model based on https://huggingface.co/sberbank-ai/sbert_large_nlu_ru',
-      author='Almira Murtazina',
-      author_email='ar.murtazina@skbkontur.ru',
-      packages=['sbert_punc_case_ru'],
-      install_requires=['transformers>=4.18.3'],
-      classifiers=[
-              "Operating System :: OS Independent",
-              "Programming Language :: Python :: 3",
-              "Programming Language :: Python :: 3.6",
-              "Programming Language :: Python :: 3.7",
-              "Programming Language :: Python :: 3.8",
-              "Programming Language :: Python :: 3.9",
-              "Topic :: Scientific/Engineering :: Artificial Intelligence",
-          ]
-     )
+setup(
+    name="sbert_punc_case_ru",
+    version="0.2",
+    description="Punctuation and Case Restoration model based on https://huggingface.co/sberbank-ai/sbert_large_nlu_ru",
+    author="Almira Murtazina",
+    author_email="ar.murtazina@skbkontur.ru",
+    packages=["sbert_punc_case_ru"],
+    install_requires=[
+        "transformers>=4.36.2",
+        "torch",
+        "numpy"
+    ],
+    classifiers=[
+        "Operating System :: OS Independent",
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.6",
+        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+        "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    ],
+)

tokenizer_config.json

@@ -1,13 +1,55 @@
 {
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "100": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "101": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "102": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "103": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": true,
   "cls_token": "[CLS]",
   "do_basic_tokenize": true,
   "do_lower_case": true,
   "mask_token": "[MASK]",
-  "name_or_path": "sberbank-ai/sbert_large_nlu_ru",
+  "model_max_length": 1000000000000000019884624838656,
   "never_split": null,
   "pad_token": "[PAD]",
   "sep_token": "[SEP]",
-  "special_tokens_map_file": null,
   "strip_accents": false,
   "tokenize_chinese_chars": true,
   "tokenizer_class": "BertTokenizer",