Commit 86b71954 authored by haileyschoelkopf's avatar haileyschoelkopf Committed by lintangsutawika
Browse files

add start of docstring to HFLM

parent 0e4f5361
...@@ -19,6 +19,13 @@ from accelerate import Accelerator ...@@ -19,6 +19,13 @@ from accelerate import Accelerator
@register_model("hf-auto") @register_model("hf-auto")
class HFLM(LM): class HFLM(LM):
"""
An abstracted Huggingface model class. Enables usage with both models of
`transformers.AutoModelForCausalLM` and `transformers.AutoModelForSeq2SeqLM` classes.
Supports data-parallel multi-GPU with HF Accelerate.
"""
AUTO_MODEL_CLASS = None AUTO_MODEL_CLASS = None
def __init__( def __init__(
self, self,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment