Create LyapunovTokenizer.py

LyapunovTokenizer.py

@@ -0,0 +1,92 @@
+from collections import OrderedDict
+import sympy as sp
+from transformers import PreTrainedTokenizer
+import json
+import os
+from huggingface_hub import upload_folder
+
+SPECIAL_WORDS = ["<s>", "</s>", "<pad>", "(", ")"]
+SPECIAL_WORDS = SPECIAL_WORDS + [f"<SPECIAL_{i}>" for i in range(10)]
+
+class LyapunovTokenizer(PreTrainedTokenizer):
+    def __init__(self):
+        self.SYMPY_OPERATORS = {
+            sp.Add: "+",
+            sp.Mul: "*",
+            sp.Pow: "^",
+            sp.exp: "exp",
+            sp.log: "ln",
+            sp.Abs: "Abs",
+            sp.sin: "sin",
+            sp.cos: "cos",
+            sp.tan: "tan",
+            sp.asin: "asin",
+            sp.acos: "acos",
+            sp.atan: "atan",
+            sp.DiracDelta: "delta0",
+        }
+
+        self.trig_ops = ["sin", "cos", "tan"]
+        self.arctrig_ops = ["asin", "acos", "atan"]
+        self.exp_ops = ["exp", "ln"]
+        self.other_ops = ["sqrt"]
+
+        op_set = {
+            "+": 2,
+            "-": 2,
+            "*": 2,
+            "/": 2,
+            "^": 2,
+            "sqrt": 1,
+            "exp": 1,
+            "ln": 1,
+            "sin": 1,
+            "cos": 1,
+            "tan": 1,
+            "asin": 1,
+            "acos": 1,
+            "atan": 1,
+            "Abs": 1,
+        }
+
+        self.int_base = 1000
+        self.max_degree = 6
+
+        self.operators_lyap = op_set
+        self.operators = self.operators_lyap
+
+        self.variables = OrderedDict({f"x{i}": sp.Symbol(f"x{i}") for i in range(2 * self.max_degree)})
+        self.constants = ["pi", "E"]
+        self.symbols = ["I", "INT+", "INT-", "FLOAT+", "FLOAT-", ".", "10^"]
+        self.elements = [str(i) for i in range(max(10, self.int_base))]
+        self.mask_symbol = ["<mask>"]
+
+        self.words = SPECIAL_WORDS + self.constants + list(self.variables.keys()) + list(self.operators.keys()) + self.symbols + self.elements + self.mask_symbol
+
+        self.vocab = {s: i for i, s in enumerate(self.words)}
+        self.inv_vocab = {v: k for k, v in self.vocab.items()}
+        super().__init__(
+            model_max_length=2048, bos_token="<s>", eos_token="</s>", unk_token="<unk>", mask_token="<mask>"
+        )
+
+    def _tokenize(self, text):
+        return text.split()
+
+    def _convert_token_to_id(self, token):
+        return self.vocab.get(token, self.unk_token_id)
+
+    def _convert_id_to_token(self, index):
+        return self.inv_vocab.get(index, self.unk_token)
+
+    def get_vocab(self):
+        return self.vocab
+
+    @property
+    def vocab_size(self):
+        return len(self.vocab)
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        vocab_file = os.path.join(save_directory, "vocab.json")
+        with open(vocab_file, "w") as f:
+            json.dump(self.vocab, f)
+        return (vocab_file,)