Upload 3 files

modeling_openmoe.py

@@ -28,9 +28,7 @@ from torch import nn
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
 from transformers.models.llama.configuration_llama import LlamaConfig
-# ========= Disable Flash Attn =============
 # from .llama_attn import LlamaAttention
-# ========= Disable Flash Attn =============
 
 from transformers.utils import (
     add_start_docstrings,
@@ -399,11 +397,7 @@ class OpenMoeAttention(nn.Module):
         value_states = repeat_kv(value_states, self.num_key_value_groups)
 
         if HAS_FLASH_ATTN and self.use_kernel:
-            # from flash_attn import flash_attn_func
-            # If we use `from flash_attn import flash_attn_func` directly, 
-            # AutoModelForCausalLM.from_pretrained will treat flash_attn as a compulsory dependency and raise error if it cannot be found.
-            # Here is a workaround to avoid the error.
-            exec("from flash_attn import flash_attn_func")  
+            from flash_attn import flash_attn_func  
 
             query_states = query_states.transpose(1, 2)
             key_states = key_states.transpose(1, 2)

tokenization_openmoe.py

@@ -0,0 +1,22 @@
+from transformers import T5Tokenizer
+from typing import List, Optional, Tuple, Union
+
+class OpenMoeTokenizer(T5Tokenizer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.padding_side = 'left'
+        self.add_bos_token = True
+        self.add_eos_token = False
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        if self.add_eos_token:
+            token_ids_0 = self._add_eos_if_not_present(token_ids_0)
+        if self.add_bos_token:
+            token_ids_0 = [self.pad_token_id] + token_ids_0
+        if token_ids_1 is None:
+            return token_ids_0
+        else:
+            token_ids_1 = self._add_eos_if_not_present(token_ids_1)
+            return token_ids_0 + token_ids_1

tokenizer_config.json

@@ -2745,8 +2745,13 @@
   "sp_model_kwargs": {},
   "spaces_between_special_tokens": false,
   "tokenizer_class": "OpenMoeTokenizer",
-  "tokenizer_file": "/home1/08125/fuzhao/.cache/huggingface/hub/models--google--umt5-small/snapshots/8c63c2b77efbf8e41206a2c8d994846cc9392360/tokenizer.json",
   "trust_remote_code": true,
   "unk_token": "<unk>",
-  "verbose": false
+  "verbose": false,
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_openmoe.OpenMoeTokenizer",
+      null
+      ]
+  }
 }