Upload 21 files

G2P_lexicon/G2P.py

@@ -74,8 +74,9 @@ class GraphemeToPhoneme:
         return pred
 
 
-dict_path = os.path.join(dirname, "my_tokenizer/bpe_512_lex.json")
-model_path = os.path.join(dirname, "models/model0.07.pt")
+dict_path = os.path.join(dirname, "my_tokenizer/bpe_256_cmu.json")
+model_path = os.path.join(dirname, "models/model_g2p.pt")
+
 
 tokenizer_g2p = Tokenizer.from_file(dict_path)
 g2p_model = TransformerBlock(config=config_g2p, tokenizer=tokenizer_g2p)

G2P_lexicon/SP.py

@@ -65,8 +65,8 @@ class Stress_Pred:
         return pred
 
 
-dict_path = os.path.join(dirname, "my_tokenizer\my_dict_256.json")
-model_path = os.path.join(dirname, "models\model_0.159.pt")
+dict_path = os.path.join(dirname, "my_tokenizer\sp_dict.json")
+model_path = os.path.join(dirname, "models\model_sp.pt")
 
 tokenizer_sp = Tokenizer_sp(dict_path=dict_path)
 

G2P_lexicon/config_models.py

@@ -4,12 +4,14 @@ config_sp = {
     "NUM": 3,
     "NUM_HEADS": 4,
     "MAX_LEN": 32,
+    "BIAS": True
     }
 
 config_g2p = {
-        "D_MODEL":  512,
-        "D_FF": 2048,
-        "NUM": 6,
-        "NUM_HEADS": 8,
+        "D_MODEL":  256,
+        "D_FF": 1024,
+        "NUM": 3,
+        "NUM_HEADS": 4,
         "MAX_LEN": 32,
+        "BIAS": False,
         }

G2P_lexicon/data_preparation.py

@@ -1,43 +1,59 @@
 import re
 
+one = ["", "one ", "two ", "three ", "four ",
+       "five ", "six ", "seven ", "eight ",
+       "nine ", "ten ", "eleven ", "twelve ",
+       "thirteen ", "fourteen ", "fifteen ",
+       "sixteen ", "seventeen ", "eighteen ",
+       "nineteen "]
 
-def intToWord(number):
-    ones = ("", "one", "two", "three", "four", "five", "six", "seven", "eight", "nine")
-    tens = ("", "", "twenty", "thirty", "forty", "fifty", "sixty", "seventy", "eighty", "ninety")
-    teens = (
-        "ten", "eleven", "twelve", "thirteen", "fourteen", "fifteen", "sixteen", "seventeen", "eighteen", "nineteen")
-    levels = (
-        "", "thousand", "million", "billion", "trillion", "quadrillion", "quintillion", "sextillion", "septillion",
-        "octillion", "nonillion")
-
-    word = ""
-    num = reversed(str(number))
-    number = ""
-    for x in num:
-        number += x
-    del num
-    if len(number) % 3 == 1: number += "0"
-    x = 0
-    for digit in number:
-        if x % 3 == 0:
-            word = levels[x // 3] + " " + word
-            n = int(digit)
-        elif x % 3 == 1:
-            if digit == "1":
-                num = teens[n]
-            else:
-                num = tens[int(digit)]
-                if n:
-                    if num:
-                        num +=  ones[n]
-                    else:
-                        num = ones[n]
-            word = num + " " + word
-        elif x % 3 == 2:
-            if digit != "0":
-                word = ones[int(digit)] + " hundred " + word
-        x += 1
-    return word.strip(" ")
+# strings at index 0 and 1 are not used,
+# they are to make array indexing simple
+ten = ["", "", "twenty ", "thirty ", "forty ",
+       "fifty ", "sixty ", "seventy ", "eighty ",
+       "ninety "]
+
+
+def numToWords(n, s):
+    str = ""
+
+    if n <= 19:
+        str += one[n]
+    # if n is more than 19, divide it
+    else:
+        str += ten[n // 10] + one[n % 10]
+
+    # if n is non-zero
+    if (n):
+        str += s
+
+    return str
+
+
+def intToWord(n):
+    n=int(n)
+    out = ""
+
+    out += numToWords((n // 10000000),
+                      "crore ")
+
+    out += numToWords(((n // 100000) % 100),
+                      "lakh ")
+
+    out += numToWords(((n // 1000) % 100),
+                      "thousand ")
+
+    out += numToWords(((n // 100) % 10),
+                      "hundred ")
+
+    if n > 100 and n % 100:
+        out += "and "
+
+    # handles digits at ones and tens
+    # places (if any)
+    out += numToWords((n % 100), "")
+
+    return out.strip()
 
 
 def preprocess_text(text):
@@ -48,7 +64,7 @@ def preprocess_text(text):
     return:
         ['HELLO', ',', 'WORLD', 'THIS', 'IS', 'A', 'SAMPLE', 'TEXT', 'WITH', 'NUMBERS', 'AND', 'SYMBOLS', '.']
     """
-    if not(text.isspace()) and text and text:
+    if not (text.isspace()) and text and text:
 
         text = text.upper()
         text = re.sub(r'([.,])', r' \1 ', text)

G2P_lexicon/models/model_g2p.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:07c75f15750171f0c1be7be681b433031fe9beaa1d223054cb06fd5ebfcc0fcf
+size 22952698

G2P_lexicon/models/model_sp.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ce2f8269e96abaf00086f4c61043046656deb8cf397ce7f1501d2f354dd6bea7
+size 22471914

G2P_lexicon/my_tokenizer/bpe_256_cmu.json

@@ -0,0 +1,530 @@
+{
+  "version": "1.0",
+  "truncation": null,
+  "padding": null,
+  "added_tokens": [
+    {
+      "id": 0,
+      "content": "<unk>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 256,
+      "content": "<pad>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 257,
+      "content": "<bos>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 258,
+      "content": "<eos>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    }
+  ],
+  "normalizer": {
+    "type": "BertNormalizer",
+    "clean_text": true,
+    "handle_chinese_chars": true,
+    "strip_accents": null,
+    "lowercase": false
+  },
+  "pre_tokenizer": {
+    "type": "BertPreTokenizer"
+  },
+  "post_processor": null,
+  "decoder": {
+    "type": "BPEDecoder",
+    "suffix": "</w>"
+  },
+  "model": {
+    "type": "BPE",
+    "dropout": null,
+    "unk_token": "<unk>",
+    "continuing_subword_prefix": null,
+    "end_of_word_suffix": "</w>",
+    "fuse_unk": false,
+    "byte_fallback": false,
+    "ignore_merges": false,
+    "vocab": {
+      "<unk>": 0,
+      "A": 1,
+      "B": 2,
+      "C": 3,
+      "D": 4,
+      "E": 5,
+      "F": 6,
+      "G": 7,
+      "H": 8,
+      "I": 9,
+      "J": 10,
+      "K": 11,
+      "L": 12,
+      "M": 13,
+      "N": 14,
+      "O": 15,
+      "P": 16,
+      "Q": 17,
+      "R": 18,
+      "S": 19,
+      "T": 20,
+      "U": 21,
+      "V": 22,
+      "W": 23,
+      "X": 24,
+      "Y": 25,
+      "Z": 26,
+      "Ġ": 27,
+      "Z</w>": 28,
+      "R</w>": 29,
+      "G</w>": 30,
+      "N</w>": 31,
+      "S</w>": 32,
+      "H</w>": 33,
+      "D</w>": 34,
+      "M</w>": 35,
+      "W</w>": 36,
+      "L</w>": 37,
+      "Y</w>": 38,
+      "E</w>": 39,
+      "T</w>": 40,
+      "K</w>": 41,
+      "V</w>": 42,
+      "A</w>": 43,
+      "F</w>": 44,
+      "Q</w>": 45,
+      "B</w>": 46,
+      "O</w>": 47,
+      "P</w>": 48,
+      "I</w>": 49,
+      "C</w>": 50,
+      "U</w>": 51,
+      "X</w>": 52,
+      "J</w>": 53,
+      "HĠ": 54,
+      "ĠA": 55,
+      "ĠI": 56,
+      "ĠAHĠ": 57,
+      "ĠE": 58,
+      "ĠIHĠ": 59,
+      "YĠ": 60,
+      "RĠ": 61,
+      "NĠ": 62,
+      "AĠ": 63,
+      "WĠ": 64,
+      "EĠ": 65,
+      "ĠAAĠ": 66,
+      "SĠ": 67,
+      "ĠEHĠ": 68,
+      "ĠAEĠ": 69,
+      "ĠR": 70,
+      "ĠIYĠ": 71,
+      "LĠ": 72,
+      "ĠERĠ": 73,
+      "HĠA": 74,
+      "KĠ": 75,
+      "OWĠ": 76,
+      "ĠIY</w>": 77,
+      "ĠEYĠ": 78,
+      "TĠ": 79,
+      "ĠAO": 80,
+      "GĠ": 81,
+      "UWĠ": 82,
+      "ĠAHĠNĠ": 83,
+      "ĠAOĠ": 84,
+      "ĠIHĠN": 85,
+      "IHĠ": 86,
+      "MĠ": 87,
+      "ĠAH</w>": 88,
+      "ĠAYĠ": 89,
+      "DĠ": 90,
+      "SĠT": 91,
+      "HĠE": 92,
+      "HĠAHĠ": 93,
+      "ĠAHĠN</w>": 94,
+      "ĠIY": 95,
+      "ĠER</w>": 96,
+      "PĠ": 97,
+      "BĠ": 98,
+      "AHĠ": 99,
+      "ĠIHĠNG</w>": 100,
+      "LĠAHĠ": 101,
+      "NĠAHĠ": 102,
+      "ĠER": 103,
+      "OW</w>": 104,
+      "KĠAHĠ": 105,
+      "ĠAAĠRĠ": 106,
+      "HHĠA": 107,
+      "LĠIY</w>": 108,
+      "LĠIHĠ": 109,
+      "TĠS</w>": 110,
+      "HĠIHĠ": 111,
+      "SĠIHĠ": 112,
+      "DĠIHĠ": 113,
+      "TĠIHĠ": 114,
+      "ĠAOĠRĠ": 115,
+      "ĠERĠZ</w>": 116,
+      "SĠAHĠ": 117,
+      "ĠIYĠZ</w>": 118,
+      "FĠ": 119,
+      "IN": 120,
+      "SHĠAHĠ": 121,
+      "TĠAHĠ": 122,
+      "NĠZ</w>": 123,
+      "ER": 124,
+      "AEĠ": 125,
+      "MĠAHĠ": 126,
+      "ĠAEĠNĠ": 127,
+      "HĠEHĠ": 128,
+      "EHĠ": 129,
+      "UHĠ": 130,
+      "ĠRĠAHĠ": 131,
+      "ĠAHĠNĠZ</w>": 132,
+      "BĠAHĠ": 133,
+      "ĠEHĠR": 134,
+      "ĠEHĠNĠ": 135,
+      "DĠAHĠ": 136,
+      "ĠRĠIHĠ": 137,
+      "HĠI": 138,
+      "KĠAAĠ": 139,
+      "LĠZ</w>": 140,
+      "ĠIHĠNGĠ": 141,
+      "NGĠ": 142,
+      "NĠIHĠ": 143,
+      "MĠIHĠ": 144,
+      "AN": 145,
+      "WĠIHĠ": 146,
+      "ĠAWĠ": 147,
+      "AR": 148,
+      "ZĠ": 149,
+      "AAĠ": 150,
+      "SĠT</w>": 151,
+      "YĠUWĠ": 152,
+      "DĠZ</w>": 153,
+      "RĠOWĠ": 154,
+      "AHĠNĠ": 155,
+      "SĠK": 156,
+      "EN": 157,
+      "OĠ": 158,
+      "SĠP": 159,
+      "BĠERĠ": 160,
+      "LĠAEĠ": 161,
+      "KĠS</w>": 162,
+      "RĠIHĠ": 163,
+      "IHĠNĠ": 164,
+      "TĠR": 165,
+      "ĠIYĠAHĠ": 166,
+      "ĠAAĠNĠ": 167,
+      "ON": 168,
+      "YĠAHĠ": 169,
+      "PĠAHĠ": 170,
+      "VĠ": 171,
+      "RĠAHĠ": 172,
+      "VĠIHĠ": 173,
+      "LĠEHĠ": 174,
+      "KĠAEĠ": 175,
+      "HHĠ": 176,
+      "LĠIYĠ": 177,
+      "OR": 178,
+      "HĠERĠ": 179,
+      "GĠAHĠ": 180,
+      "MĠAEĠ": 181,
+      "GĠR": 182,
+      "ST": 183,
+      "AT": 184,
+      "ES</w>": 185,
+      "BĠR": 186,
+      "RĠIYĠ": 187,
+      "BĠIHĠ": 188,
+      "SHĠ": 189,
+      "LĠEYĠ": 190,
+      "PĠR": 191,
+      "LĠAAĠ": 192,
+      "AL": 193,
+      "TĠIY</w>": 194,
+      "HHĠAEĠ": 195,
+      "SĠEHĠ": 196,
+      "NĠAHĠS</w>": 197,
+      "TH</w>": 198,
+      "EL": 199,
+      "HĠIYĠ": 200,
+      "FĠAHĠ": 201,
+      "LĠAYĠ": 202,
+      "LĠD</w>": 203,
+      "KĠW": 204,
+      "MĠEHĠ": 205,
+      "RE": 206,
+      "PĠIHĠ": 207,
+      "FĠIHĠ": 208,
+      "SHĠAHĠN</w>": 209,
+      "NĠIY</w>": 210,
+      "MĠAAĠ": 211,
+      "KĠR": 212,
+      "VĠAHĠ": 213,
+      "THĠ": 214,
+      "UW</w>": 215,
+      "OWĠZ</w>": 216,
+      "HHĠAAĠ": 217,
+      "CH": 218,
+      "RĠUWĠ": 219,
+      "OYĠ": 220,
+      "ĠAOĠR": 221,
+      "KĠIHĠ": 222,
+      "HĠAEĠ": 223,
+      "ED</w>": 224,
+      "ZĠAHĠ": 225,
+      "HHĠEHĠ": 226,
+      "SĠIHĠZ</w>": 227,
+      "DĠEHĠ": 228,
+      "JHĠAHĠ": 229,
+      "JHĠIHĠ": 230,
+      "BĠAEĠ": 231,
+      "TĠERĠ": 232,
+      "JHĠ": 233,
+      "OW": 234,
+      "BĠEHĠ": 235,
+      "SĠIYĠ": 236,
+      "OWĠLĠ": 237,
+      "VĠERĠ": 238,
+      "ĠEY</w>": 239,
+      "TĠIHĠD</w>": 240,
+      "KĠAHĠNĠ": 241,
+      "LE": 242,
+      "MĠAHĠN</w>": 243,
+      "ĠAHĠNĠT</w>": 244,
+      "RĠEHĠ": 245,
+      "NĠAH</w>": 246,
+      "CHĠ": 247,
+      "IS": 248,
+      "UW": 249,
+      "PĠERĠ": 250,
+      "SĠTĠ": 251,
+      "PĠAAĠ": 252,
+      "TĠAHĠN</w>": 253,
+      "LĠUWĠ": 254,
+      "HĠAAĠ": 255
+    },
+    "merges": [
+      "H Ġ",
+      "Ġ A",
+      "Ġ I",
+      "ĠA HĠ",
+      "Ġ E",
+      "ĠI HĠ",
+      "Y Ġ",
+      "R Ġ",
+      "N Ġ",
+      "A Ġ",
+      "W Ġ",
+      "E Ġ",
+      "ĠA AĠ",
+      "S Ġ",
+      "ĠE HĠ",
+      "ĠA EĠ",
+      "Ġ R",
+      "ĠI YĠ",
+      "L Ġ",
+      "ĠE RĠ",
+      "HĠ A",
+      "K Ġ",
+      "O WĠ",
+      "ĠI Y</w>",
+      "ĠE YĠ",
+      "T Ġ",
+      "ĠA O",
+      "G Ġ",
+      "U WĠ",
+      "ĠAHĠ NĠ",
+      "ĠAO Ġ",
+      "ĠIHĠ N",
+      "I HĠ",
+      "M Ġ",
+      "ĠA H</w>",
+      "ĠA YĠ",
+      "D Ġ",
+      "SĠ T",
+      "HĠ E",
+      "HĠA HĠ",
+      "ĠAHĠ N</w>",
+      "ĠI Y",
+      "ĠE R</w>",
+      "P Ġ",
+      "B Ġ",
+      "A HĠ",
+      "ĠIHĠN G</w>",
+      "L ĠAHĠ",
+      "N ĠAHĠ",
+      "ĠE R",
+      "O W</w>",
+      "K ĠAHĠ",
+      "ĠAAĠ RĠ",
+      "H HĠA",
+      "L ĠIY</w>",
+      "L ĠIHĠ",
+      "TĠ S</w>",
+      "HĠ IHĠ",
+      "S ĠIHĠ",
+      "D ĠIHĠ",
+      "T ĠIHĠ",
+      "ĠAOĠ RĠ",
+      "ĠERĠ Z</w>",
+      "S ĠAHĠ",
+      "ĠIYĠ Z</w>",
+      "F Ġ",
+      "I N",
+      "S HĠAHĠ",
+      "T ĠAHĠ",
+      "NĠ Z</w>",
+      "E R",
+      "A EĠ",
+      "M ĠAHĠ",
+      "ĠAEĠ NĠ",
+      "HĠE HĠ",
+      "E HĠ",
+      "U HĠ",
+      "ĠR ĠAHĠ",
+      "ĠAHĠNĠ Z</w>",
+      "B ĠAHĠ",
+      "ĠEHĠ R",
+      "ĠEHĠ NĠ",
+      "D ĠAHĠ",
+      "ĠR ĠIHĠ",
+      "HĠ I",
+      "K ĠAAĠ",
+      "LĠ Z</w>",
+      "ĠIHĠN GĠ",
+      "N GĠ",
+      "N ĠIHĠ",
+      "M ĠIHĠ",
+      "A N",
+      "W ĠIHĠ",
+      "ĠA WĠ",
+      "A R",
+      "Z Ġ",
+      "A AĠ",
+      "SĠ T</w>",
+      "YĠ UWĠ",
+      "DĠ Z</w>",
+      "RĠ OWĠ",
+      "AHĠ NĠ",
+      "SĠ K",
+      "E N",
+      "O Ġ",
+      "SĠ P",
+      "B ĠERĠ",
+      "L ĠAEĠ",
+      "KĠ S</w>",
+      "R ĠIHĠ",
+      "IHĠ NĠ",
+      "T ĠR",
+      "ĠIY ĠAHĠ",
+      "ĠAAĠ NĠ",
+      "O N",
+      "Y ĠAHĠ",
+      "P ĠAHĠ",
+      "V Ġ",
+      "R ĠAHĠ",
+      "V ĠIHĠ",
+      "L ĠEHĠ",
+      "K ĠAEĠ",
+      "H HĠ",
+      "L ĠIYĠ",
+      "O R",
+      "HĠE RĠ",
+      "G ĠAHĠ",
+      "M ĠAEĠ",
+      "G ĠR",
+      "S T",
+      "A T",
+      "E S</w>",
+      "B ĠR",
+      "R ĠIYĠ",
+      "B ĠIHĠ",
+      "S HĠ",
+      "L ĠEYĠ",
+      "P ĠR",
+      "L ĠAAĠ",
+      "A L",
+      "T ĠIY</w>",
+      "HHĠA EĠ",
+      "S ĠEHĠ",
+      "NĠAHĠ S</w>",
+      "T H</w>",
+      "E L",
+      "HĠI YĠ",
+      "F ĠAHĠ",
+      "L ĠAYĠ",
+      "LĠ D</w>",
+      "KĠ W",
+      "M ĠEHĠ",
+      "R E",
+      "P ĠIHĠ",
+      "F ĠIHĠ",
+      "SHĠAHĠ N</w>",
+      "N ĠIY</w>",
+      "M ĠAAĠ",
+      "K ĠR",
+      "V ĠAHĠ",
+      "T HĠ",
+      "U W</w>",
+      "OWĠ Z</w>",
+      "HHĠA AĠ",
+      "C H",
+      "RĠ UWĠ",
+      "O YĠ",
+      "ĠAO ĠR",
+      "K ĠIHĠ",
+      "HĠA EĠ",
+      "E D</w>",
+      "Z ĠAHĠ",
+      "H HĠEHĠ",
+      "SĠIHĠ Z</w>",
+      "D ĠEHĠ",
+      "J HĠAHĠ",
+      "J HĠIHĠ",
+      "B ĠAEĠ",
+      "T ĠERĠ",
+      "J HĠ",
+      "O W",
+      "B ĠEHĠ",
+      "S ĠIYĠ",
+      "OWĠ LĠ",
+      "V ĠERĠ",
+      "ĠE Y</w>",
+      "TĠIHĠ D</w>",
+      "K ĠAHĠNĠ",
+      "L E",
+      "M ĠAHĠN</w>",
+      "ĠAHĠNĠ T</w>",
+      "R ĠEHĠ",
+      "N ĠAH</w>",
+      "C HĠ",
+      "I S",
+      "U W",
+      "P ĠERĠ",
+      "SĠ TĠ",
+      "P ĠAAĠ",
+      "T ĠAHĠN</w>",
+      "LĠ UWĠ",
+      "HĠA AĠ"
+    ]
+  }
+}

G2P_lexicon/my_tokenizer/sp_dict.json

@@ -0,0 +1,90 @@
+{
+  "0": "<sos>",
+  "1": "<eos>",
+  "2": "<unk>",
+  "3": "<pad>",
+  "4": "AA1",
+  "5": "UW",
+  "6": "ER0",
+  "7": "F",
+  "8": "CH",
+  "9": "S",
+  "10": "AO1",
+  "11": "DH",
+  "12": "TH",
+  "13": "IY",
+  "14": "OW",
+  "15": "AH2",
+  "16": "W",
+  "17": "AH1",
+  "18": "AO",
+  "19": "D",
+  "20": "AW1",
+  "21": "OY2",
+  "22": "AO0",
+  "23": "EY0",
+  "24": "AH",
+  "25": "AE",
+  "26": "UH2",
+  "27": "OW2",
+  "28": "UW0",
+  "29": "UW1",
+  "30": "UH1",
+  "31": "ER",
+  "32": "EH2",
+  "33": "UW2",
+  "34": "ER2",
+  "35": "OY",
+  "36": "AE0",
+  "37": "AY",
+  "38": "K",
+  "39": "AA0",
+  "40": "T",
+  "41": "EH0",
+  "42": "SH",
+  "43": "ER1",
+  "44": "G",
+  "45": "EY",
+  "46": "AH0",
+  "47": "IH0",
+  "48": "L",
+  "49": "AE2",
+  "50": "B",
+  "51": "OY0",
+  "52": "EH",
+  "53": "AA2",
+  "54": "IH",
+  "55": "M",
+  "56": "AY0",
+  "57": "UH",
+  "58": "EY2",
+  "59": "IY2",
+  "60": "EY1",
+  "61": "HH",
+  "62": "P",
+  "63": "AE1",
+  "64": "OW1",
+  "65": "R",
+  "66": "IH1",
+  "67": "Z",
+  "68": "IH2",
+  "69": "IY0",
+  "70": "V",
+  "71": "JH",
+  "72": "OY1",
+  "73": "Y",
+  "74": "N",
+  "75": "AO2",
+  "76": "AW",
+  "77": "UH0",
+  "78": "IY1",
+  "79": "AW0",
+  "80": "AA",
+  "81": "NG",
+  "82": "AY1",
+  "83": "EH1",
+  "84": "AY2",
+  "85": "OW0",
+  "86": "AW2",
+  "87": "ZH"
+}

G2P_lexicon/sp_tokenizer.py

@@ -83,5 +83,5 @@ class Tokenizer_sp:
 
 
 if __name__ == "__main__":
-    tokenizer_sp = Tokenizer_sp(dict_path='./my_tokenizer/my_dict_256.json')
+    tokenizer_sp = Tokenizer_sp(dict_path='my_tokenizer/sp_dict.json')
     print(tokenizer_sp.idx2token)

G2P_lexicon/transformer.py

@@ -22,7 +22,7 @@ class PositionalEncoding(nn.Module):
 
 
 class MultiHeadSelfAttention(nn.Module):
-    def __init__(self, d_model, num_heads):
+    def __init__(self, d_model, num_heads, bias=False):
         super(MultiHeadSelfAttention, self).__init__()
         assert d_model % num_heads == 0, "d_model must be divisible by num_heads"
 
@@ -30,9 +30,9 @@ class MultiHeadSelfAttention(nn.Module):
         self.num_heads = num_heads
         self.depth = d_model // num_heads
 
-        self.wq = nn.Linear(d_model, d_model)
-        self.wk = nn.Linear(d_model, d_model)
-        self.wv = nn.Linear(d_model, d_model)
+        self.wq = nn.Linear(d_model, d_model, bias)
+        self.wk = nn.Linear(d_model, d_model, bias)
+        self.wv = nn.Linear(d_model, d_model, bias)
 
         self.fc = nn.Linear(d_model, d_model)
 
@@ -76,9 +76,9 @@ class FeedForwardNetwork(nn.Module):
 
 
 class EncoderLayer(nn.Module):
-    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
+    def __init__(self, d_model, num_heads, d_ff, dropout=0.1, bias=False):
         super(EncoderLayer, self).__init__()
-        self.self_attn = MultiHeadSelfAttention(d_model, num_heads)
+        self.self_attn = MultiHeadSelfAttention(d_model, num_heads, bias)
         self.ffn = FeedForwardNetwork(d_model, d_ff, dropout)
 
         self.layernorm1 = nn.LayerNorm(d_model)
@@ -95,10 +95,10 @@ class EncoderLayer(nn.Module):
 
 
 class DecoderLayer(nn.Module):
-    def __init__(self, d_model, num_heads, d_ff, dropout=0.1):
+    def __init__(self, d_model, num_heads, d_ff, dropout=0.1, bias=False):
         super(DecoderLayer, self).__init__()
-        self.self_attn = MultiHeadSelfAttention(d_model, num_heads)
-        self.cross_attn = MultiHeadSelfAttention(d_model, num_heads)
+        self.self_attn = MultiHeadSelfAttention(d_model, num_heads, bias)
+        self.cross_attn = MultiHeadSelfAttention(d_model, num_heads, bias)
         self.ffn = FeedForwardNetwork(d_model, d_ff, dropout)
 
         self.layernorm1 = nn.LayerNorm(d_model)
@@ -132,6 +132,7 @@ class TransformerBlock(nn.Module):
         self.num_decoder_layers = config.get('NUM', 6)
         self.d_ff = config.get('D_FF', 2048)
         self.dropout = config.get('DROPOUT', 0.1)
+        self.bias = config.get('BIAS', False)
         self.stress = stress
 
         self.encoder_embedding = nn.Embedding(self.input_vocab_size, self.d_model)
@@ -140,10 +141,10 @@ class TransformerBlock(nn.Module):
         self.pos_embedding = PositionalEncoding(self.d_model, config.get('MAX_LEN', 32))
 
         self.encoder_layers = nn.ModuleList(
-            [EncoderLayer(self.d_model, self.num_heads, self.d_ff, self.dropout) for _ in
+            [EncoderLayer(self.d_model, self.num_heads, self.d_ff, self.dropout, self.bias) for _ in
              range(self.num_encoder_layers)])
         self.decoder_layers = nn.ModuleList(
-            [DecoderLayer(self.d_model, self.num_heads, self.d_ff, self.dropout) for _ in
+            [DecoderLayer(self.d_model, self.num_heads, self.d_ff, self.dropout, self.bias) for _ in
              range(self.num_decoder_layers)])
 
         self.fc_out = nn.Linear(self.d_model, self.target_vocab_size)