fix #13

hiyouga · Jul 19, 2023 · 068ec04 · 068ec04
1 parent 4091ad5
commit 068ec04
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -16,6 +16,7 @@ This repo aims to assist the developers with injecting **fresh** and **customize
 
 - [GPT-J](https://huggingface.co/EleutherAI/gpt-j-6b) (6B)
 - [LLaMA](https://github.com/facebookresearch/llama) (7B/13B)
+- [LLaMA-2](https://huggingface.co/meta-llama) (7B/13B)
 - [BLOOM](https://huggingface.co/bigscience/bloomz) (7.1B)
 - [Falcon](https://huggingface.co/tiiuae/falcon-7b) (7B)
 - [Baichuan](https://huggingface.co/baichuan-inc/Baichuan-7B) (7B/13B)

diff --git a/fastedit/rome/rome_hparams.py b/fastedit/rome/rome_hparams.py
@@ -53,11 +53,11 @@ def from_name(cls, name: str):
             mom2_dtype="float16"
         )
 
-        if name == "gpj-j-6b":
+        if name == "gpt-j-6b":
             pass
         elif name == "llama-7b":
             r"""
-            Supports: LLaMA-7B, Baichuan-7B, InternLM-7B...
+            Supports: LLaMA-7B, LLaMA-2-7B, Baichuan-7B, InternLM-7B...
             """
             data.update(dict(
                 v_loss_layer=31,
@@ -69,7 +69,7 @@ def from_name(cls, name: str):
             ))
         elif name == "llama-13b":
             r"""
-            Supports LLaMA-13B, Baichuan-13B...
+            Supports LLaMA-13B, LLaMA-2-13B, Baichuan-13B...
             """
             data.update(dict(
                 layers=[10],