flax model

huggingface · Dec 28, 2021 · 52fb413 · 52fb413
1 parent dc90de9
commit 52fb413
Show file tree

Hide file tree

Showing 4 changed files with 399 additions and 1,231 deletions.
diff --git a/src/transformers/__init__.py b/src/transformers/__init__.py
@@ -2110,6 +2110,7 @@
         [
             "FlaxXGLMModel",
             "FlaxXGLMPreTrainedModel",
+            "FlaxXGLMForCausalLM",
         ]
     )
 else:
@@ -3847,7 +3848,7 @@
             FlaxWav2Vec2Model,
             FlaxWav2Vec2PreTrainedModel,
         )
-        from .models.xglm import FlaxXGLMModel, FlaxXGLMPreTrainedModel
+        from .models.xglm import FlaxXGLMForCausalLM, FlaxXGLMModel, FlaxXGLMPreTrainedModel
     else:
         # Import the same objects as dummies to get them in the namespace.
         # They will raise an import error if the user tries to instantiate / use them.

diff --git a/src/transformers/models/xglm/__init__.py b/src/transformers/models/xglm/__init__.py
@@ -42,6 +42,7 @@
     _import_structure["modeling_flax_xglm"] = [
         "FlaxXGLMModel",
         "FlaxXGLMPreTrainedModel",
+        "FlaxXGLMForCausalLM",
     ]
 
 
@@ -56,7 +57,7 @@
         from .modeling_xglm import XGLM_PRETRAINED_MODEL_ARCHIVE_LIST, XGLMForCausalLM, XGLMModel, XGLMPreTrainedModel
 
     if is_flax_available():
-        from .modeling_xglm import FlaxXGLMModel, FlaxXGLMPreTrainedModel
+        from .modeling_xglm import FlaxXGLMForCausalLM, FlaxXGLMModel, FlaxXGLMPreTrainedModel
 
 
 else: