some formatting

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",
+    ],
+)