baidu
/

ERNIE-4.5-300B-A47B-Base-PT

Text Generation

Model card Files Files and versions

AntonV HF Staff commited on Jul 21

Commit

36ef824

·

1 Parent(s): f952b54

readme

Files changed (1) hide show

README.md +6 -1

README.md CHANGED Viewed

@@ -79,7 +79,12 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 model_name = "baidu/ERNIE-4.5-300B-A47B-Base-PT"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
 prompt = "Large language model is"
 model_inputs = tokenizer([prompt], add_special_tokens=False, return_tensors="pt").to(model.device)

 model_name = "baidu/ERNIE-4.5-300B-A47B-Base-PT"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+    tp_plan="auto",
+)
 prompt = "Large language model is"
 model_inputs = tokenizer([prompt], add_special_tokens=False, return_tensors="pt").to(model.device)