Update model

config.json

@@ -0,0 +1,66 @@
+{
+  "architectures": [
+    "GreedyModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_greedy.GreedyConfig",
+    "AutoModelForCausalLM": "modeling_greedy.GreedyModel"
+  },
+  "eos_token_id": 8,
+  "pad_token_id": 9,
+  "reciprocals": [
+    [
+      4,
+      3
+    ],
+    [
+      5,
+      2
+    ],
+    [
+      6,
+      1
+    ]
+  ],
+  "reducables": [
+    [
+      [
+        4
+      ],
+      [
+        3
+      ]
+    ],
+    [
+      [
+        5
+      ],
+      [
+        2
+      ]
+    ],
+    [
+      [
+        6
+      ],
+      [
+        1
+      ]
+    ],
+    [
+      [
+        4,
+        5,
+        6
+      ],
+      [
+        1,
+        2,
+        3
+      ]
+    ]
+  ],
+  "torch_dtype": "float32",
+  "transformers_version": "4.21.1",
+  "vocab_size": 10
+}

configuration_greedy.py

@@ -0,0 +1,39 @@
+from transformers import PretrainedConfig, PreTrainedTokenizerBase
+from freegroup import tools
+
+class GreedyConfig(PretrainedConfig):
+
+    @classmethod
+    def from_tokenizer(cls, freegroup_dimension, tokenizer: PreTrainedTokenizerBase, **kwargs):
+
+        freegroup_generators = list(range(1, freegroup_dimension + 1))
+
+        reciprocals = []
+        for x in freegroup_generators:
+            a, b = tokenizer.convert_tokens_to_ids([str(x), str(-x)])
+            reciprocals.append([a, b])
+
+        reducables = [[] for _ in range(freegroup_dimension + 1)]
+        for reducable, closure_generator in zip(reducables, [[x] for x in freegroup_generators] + [freegroup_generators[::]]):
+            reducable.append(tokenizer.convert_tokens_to_ids(list(map(str, closure_generator))))
+            reducable.append(tokenizer.convert_tokens_to_ids(list(map(str, tools.reciprocal(closure_generator)))))
+        
+        return cls(
+            reciprocals = reciprocals,
+            reducables = reducables,
+            vocab_size = len(tokenizer),
+            eos_token_id = tokenizer.eos_token_id,
+            pad_token_id = tokenizer.pad_token_id,
+            **kwargs
+        )
+
+    def __init__(self, **kwargs):
+        # reciporcals: List[List[int]]: i.e. ['x', 'X'], ...
+        self.reciprocals = kwargs.pop('reciprocals', None)
+
+        # reducables: List[List[List[int]]]: generators for normal closures, i.e [[[x], [X]], [[y], [Y]], ...]
+        self.reducables = kwargs.pop('reducables', None)
+
+        super().__init__(**kwargs)
+
+

modeling_greedy.py

@@ -0,0 +1,85 @@
+import torch
+
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_greedy import GreedyConfig
+from freegroup import tools
+
+class GreedyModel(PreTrainedModel):
+    config_class = GreedyConfig
+    
+    def __init__(self, config: GreedyConfig): 
+        super().__init__(config)
+        self.stub = torch.nn.parameter.Parameter(torch.tensor(0.))
+
+    def _reduce_step(self, token, stack, reducables):
+        stack.append(token.item())
+
+        for reducable in self.config.reciprocals + reducables:
+            n = len(reducable)
+            if len(stack) >= len(reducable):
+                if tools.occurs(stack[-n:], reducable * 2):
+                    del stack[-n:]
+
+        return stack
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        past = kwargs.pop('past', None)
+        return {'input_ids': input_ids, 'past': past}
+
+    def forward(self, input_ids = None, past = None, **kwargs):
+
+        assert (input_ids is not None), "Can't be None"
+        
+        batch_size, sequence_length = input_ids.shape
+
+        if past is None:
+            stacks = [[[] for _ in range(len(self.config.reducables))] for _ in range(batch_size)]
+            hidden_states = None
+        else:
+            stacks, hidden_states = past
+
+        begin_idx = 0 if hidden_states is None else hidden_states.size(0)
+
+        for t in range(begin_idx, sequence_length):
+            last_hidden_states = torch.zeros((batch_size, self.config.vocab_size))
+            
+            for batch_idx, word in enumerate(input_ids):
+                for stack, reducables in zip(stacks[batch_idx], self.config.reducables):
+                    
+                    self._reduce_step(word[t], stack, reducables)
+                    if not stack: continue
+                    
+                    last = stack[-1]
+
+                    for r in reducables:
+                        if not last in r:
+                            key = r[0]
+                            last_hidden_states[batch_idx][r[0]] += 1
+                        if last in r:
+                            pos = r.index(last)
+                            key = r[(pos + 1) % len(r)]
+                            last_hidden_states[batch_idx][key] += 1
+                    for r in self.config.reciprocals:
+                        if last in r:
+                            pos = r.index(last)
+                            key = r[(pos + 1) % len(r)]
+                            last_hidden_states[batch_idx][key] += 1
+
+                for r in self.config.reciprocals:
+                    if word[t] in r:
+                        pos = r.index(word[t])
+                        key = r[(pos + 1) % len(r)]
+                        last_hidden_states[batch_idx][key] = -torch.inf
+            
+                if all(map(lambda x: len(x) == 0, stacks[batch_idx])):
+                    last_hidden_states[batch_idx][self.config.eos_token_id] = torch.inf
+
+            if hidden_states is None: hidden_states = last_hidden_states.clone().unsqueeze(0)
+            else: hidden_states = torch.cat((hidden_states, last_hidden_states.unsqueeze(0)))
+
+        return CausalLMOutputWithPast(
+            logits = hidden_states.permute(1, 0, 2),
+            past_key_values = (stacks, hidden_states)
+        )

pytorch_model.bin

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