SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

Templates overhaul 1 (#8993)

Browse Source
This commit is contained in:
Lysandre Debut
2020-12-08 18:00:07 -05:00
committed by GitHub
parent 447808c85f
commit 67ff1c314a
16 changed files with 759 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/transformers/file_utils.py
View File

@@ -849,7 +849,7 @@ def add_code_sample_docstrings(
elif "MaskedLM" in model_class or model_class in ["FlaubertWithLMHeadModel", "XLMWithLMHeadModel"]:
doc_kwargs["mask"] = "[MASK]" if mask is None else mask
code_sample = TF_MASKED_LM_SAMPLE if is_tf_class else PT_MASKED_LM_SAMPLE
elif "LMHead" in model_class:
elif "LMHead" in model_class or "CausalLM" in model_class:
code_sample = TF_CAUSAL_LM_SAMPLE if is_tf_class else PT_CAUSAL_LM_SAMPLE
elif "Model" in model_class or "Encoder" in model_class:
code_sample = TF_BASE_MODEL_SAMPLE if is_tf_class else PT_BASE_MODEL_SAMPLE