Initial toxic model implementation, related to #53

bureaucratic-labs · Feb 28, 2020 · f1d97aa · f1d97aa
1 parent 660b8c0
commit f1d97aa
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 1 deletion.
diff --git a/dostoevsky/data/__init__.py b/dostoevsky/data/__init__.py
@@ -11,7 +11,8 @@
 
 AVAILABLE_FILES: typing.Dict[str, typing.Tuple[str, str]] = {
     'cnn-social-network-model': ('models/cnn-social-network-model.tar.xz', 'models/cnn-social-network-model.tar.xz'),
-    'fasttext-social-network-model': ('models/fasttext-social-network-model.tar.xz', 'models/fasttext-social-network-model.tar.xz')
+    'fasttext-social-network-model': ('models/fasttext-social-network-model.tar.xz', 'models/fasttext-social-network-model.tar.xz'),
+    'fasttext-toxic-model': ('models/fasttext-toxic-model.tar.xz', 'models/fasttext-toxic-model.tar.xz'),
 }
 
 

diff --git a/dostoevsky/models.py b/dostoevsky/models.py
@@ -40,6 +40,7 @@ def predict(self, sentences: List[str]):
     def get_raw_model(self):
         raise NotImplementedError
 
+
 class FastTextSocialNetworkModel(BaseModel):
     '''
     FastText model trained on RuSentiment dataset.
@@ -90,3 +91,17 @@ def predict(self, sentences: List[str], k: int = -1) -> List[
                 scores
             )) for labels, scores in Y
         ]
+
+
+class FastTextToxicModel(FastTextSocialNetworkModel):
+    '''
+    FastText model trained on russian toxic comments dataset.
+    '''
+
+    SENTENCE_LENGTH: Optional[int] = None
+
+    MODEL_PATH: str = os.path.join(
+        DATA_BASE_PATH,
+        'models/fasttext-toxic-model.bin'
+    )
+