Update model

model_consts.py

@@ -4,6 +4,6 @@ else:
     from .utils import get_upenn_tags_dict
 
 input_size = len(get_upenn_tags_dict())
-embedding_size = 128
-hidden_size = 128
+embedding_size = 256
+hidden_size = 256
 num_layers = 2

segmenter.ckpt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:26c09246a1ed23aa5be9656e36878d30b1b39aa649dbd9a24bbef7ecee5a4e7d
-size 2665888
+oid sha256:a8e6209584d0021684bb3a09ec1b717843f3086dfcc6411c57276f743f8e62fa
+size 10584544

train.py

@@ -26,6 +26,6 @@ if __name__ == "__main__":
     
     model.to(device)
 
-    train_bidirlstm_embedding_model(model, dataset, num_epochs=75, batch_size=2)
+    train_bidirlstm_embedding_model(model, dataset, num_epochs=150, batch_size=2)
 
     torch.save(model.state_dict(), "segmenter.ckpt")

utils.py

@@ -4,6 +4,64 @@ from stable_whisper.result import WordTiming
 import numpy as np
 import torch
 
+additional_tags = {
+    "as": "`AS",
+    "and": "`AND",
+    "of": "`OF",
+    "how": "`HOW",
+    "but": "`BUT",
+    "the": "`THE",
+    "a": "`A",
+    "an": "`A",
+    "which": "`WHICH",
+    "what": "`WHAT",
+    "where": "`WHERE",
+    "that": "`THAT",
+    "who": "`WHO",
+    "when": "`WHEN",
+}
+
+def get_upenn_tags_dict():
+    # tagger = PerceptronTagger()
+
+    # tags = list(tagger.tagdict.values())
+
+    # # https://www.ling.upenn.edu/courses/Fall_2003/ling001/penn_treebank_pos.html
+    # tags.extend(["CC", "CD", "DT", "EX", "FW", "IN", "JJ", "JJR", "JJS", "LS", "MD", "NN", "NNS", "NNP", "NNPS", "PDT", "POS", "PRP", "PRP$", "RB", "RBR", "RBS", "RP", "SYM", "TO", "UH", "VB", "VBD", "VBG", "VBN", "VBP", "VBZ", "WDT", "WP", "WP$", "WRB"])
+    # tags = list(set(tags))
+    # tags.sort()
+    # tags.append("BREAK")
+
+    # tags_dict = dict()
+
+    # for index, tag in enumerate(tags):
+    #     tags_dict[tag] = index
+
+    return {'#': 0, '$': 1, "''": 2,'(': 3,')': 4,',': 5,'.': 6,':': 7,'CC': 8,'CD': 9,'DT': 10,'EX': 11,'FW': 12,'IN': 13,'JJ': 14,'JJR': 15,'JJS': 16,'LS': 17,'MD': 18,'NN': 19,'NNP': 20,'NNPS': 21,'NNS': 22,'PDT': 23,'POS': 24,'PRP': 25,'PRP$': 26,'RB': 27,'RBR': 28,'RBS': 29,'RP': 30,'SYM': 31,'TO': 32,'UH': 33,'VB': 34,'VBD': 35,'VBG': 36,'VBN': 37,'VBP': 38,'VBZ': 39,'WDT': 40,'WP': 41,'WP$': 42,'WRB': 43,'``': 44,'BREAK': 45,
+            '`AS': 46,
+            '`AND': 47,
+            '`OF': 48,
+            '`HOW': 49,
+            '`BUT': 50,
+            '`THE': 51,
+            '`A': 52,
+            '`WHICH': 53,
+            '`WHAT': 54,
+            '`WHERE': 55,
+            '`THAT': 56,
+            '`WHO': 57,
+            '`WHEN': 58
+    }
+
+def nltk_extend_tags(tagged_text: list[tuple[str, str]]):
+    result = []
+    for text, tag in tagged_text:
+        text_lower = text.lower().strip()
+        if text_lower in additional_tags:
+            yield (text, additional_tags[text_lower])
+        else:
+            yield (text, tag)
+
 def bind_wordtimings_to_tags(wt: list[WordTiming]):
     raw_words = [w.word for w in wt]
 
@@ -16,6 +74,7 @@ def bind_wordtimings_to_tags(wt: list[WordTiming]):
         tokens_wordtiming_map.append(len(tokens_word))
     
     tagged_words = nltk.pos_tag(tokenized_raw_words)
+    tagged_words = list(nltk_extend_tags(tagged_words))
 
     grouped_tags = []
 
@@ -49,6 +108,7 @@ def tag_training_data(filename: str):
 
     tokenized_full_text = nltk.word_tokenize(full_text)
     tagged_full_text = nltk.pos_tag(tokenized_full_text)
+    tagged_full_text = list(nltk_extend_tags(tagged_full_text))
 
     tagged_full_text_copy = tagged_full_text
 
@@ -75,70 +135,6 @@ def tag_training_data(filename: str):
     
     return reconstructed_tags
 
-def get_upenn_tags_dict():
-    # tagger = PerceptronTagger()
-
-    # tags = list(tagger.tagdict.values())
-
-    # # https://www.ling.upenn.edu/courses/Fall_2003/ling001/penn_treebank_pos.html
-    # tags.extend(["CC", "CD", "DT", "EX", "FW", "IN", "JJ", "JJR", "JJS", "LS", "MD", "NN", "NNS", "NNP", "NNPS", "PDT", "POS", "PRP", "PRP$", "RB", "RBR", "RBS", "RP", "SYM", "TO", "UH", "VB", "VBD", "VBG", "VBN", "VBP", "VBZ", "WDT", "WP", "WP$", "WRB"])
-    # tags = list(set(tags))
-    # tags.sort()
-    # tags.append("BREAK")
-
-    # tags_dict = dict()
-
-    # for index, tag in enumerate(tags):
-    #     tags_dict[tag] = index
-
-    return {'#': 0,
-            '$': 1,
-            "''": 2,
-            '(': 3,
-            ')': 4,
-            ',': 5,
-            '.': 6,
-            ':': 7,
-            'CC': 8,
-            'CD': 9,
-            'DT': 10,
-            'EX': 11,
-            'FW': 12,
-            'IN': 13,
-            'JJ': 14,
-            'JJR': 15,
-            'JJS': 16,
-            'LS': 17,
-            'MD': 18,
-            'NN': 19,
-            'NNP': 20,
-            'NNPS': 21,
-            'NNS': 22,
-            'PDT': 23,
-            'POS': 24,
-            'PRP': 25,
-            'PRP$': 26,
-            'RB': 27,
-            'RBR': 28,
-            'RBS': 29,
-            'RP': 30,
-            'SYM': 31,
-            'TO': 32,
-            'UH': 33,
-            'VB': 34,
-            'VBD': 35,
-            'VBG': 36,
-            'VBN': 37,
-            'VBP': 38,
-            'VBZ': 39,
-            'WDT': 40,
-            'WP': 41,
-            'WP$': 42,
-            'WRB': 43,
-            '``': 44,
-            'BREAK': 45}
-
-
 def parse_tags(reconstructed_tags):
     """
         Parse reconstructed tags into input/tag datapoint.