fix typo

SeanLee97 · May 22, 2024 · 62f2cb1 · 62f2cb1
1 parent e47f220
commit 62f2cb1
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -183,7 +183,7 @@ angle = AnglE.from_pretrained('NousResearch/Llama-2-7b-hf',
                               pooling_strategy='last',
                               is_llm=True,
                               apply_billm=True,
-                              billm_model_class='LlamaForCausalMask',
+                              billm_model_class='LlamaForCausalLM',
                               torch_dtype='float16').cuda()
 
 doc_vecs = angle.encode([

diff --git a/docs/notes/quickstart.rst b/docs/notes/quickstart.rst
@@ -103,7 +103,7 @@ Specify `apply_billm` and `billm_model_class` to load and infer billm models
                                 pooling_strategy='last',
                                 is_llm=True,
                                 apply_billm=True,
-                                billm_model_class='LlamaForCausalMask',
+                                billm_model_class='LlamaForCausalLM',
                                 torch_dtype='float16').cuda()
 
     doc_vecs = angle.encode([