From bc068ab275bba18548ce14b1072b27eaf60be504 Mon Sep 17 00:00:00 2001
From: KerfuffleV2 <kerfliffle@keemail.me>
Date: Sat, 4 Nov 2023 12:35:21 -0600
Subject: [PATCH] gguf-py: Support 01.AI Yi models

---
 gguf-py/gguf/gguf.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/gguf-py/gguf/gguf.py b/gguf-py/gguf/gguf.py
index 727b4e554..a2271d225 100644
--- a/gguf-py/gguf/gguf.py
+++ b/gguf-py/gguf/gguf.py
@@ -393,6 +393,7 @@ class TensorNameMap:
             "layers.{bid}.attention_norm",                         # llama-pth
             "encoder.layer.{bid}.attention.output.LayerNorm",      # bert
             "language_model.encoder.layers.{bid}.input_layernorm", # persimmon
+            "model.layers.{bid}.ln1",                              # yi
         ),
 
         # Attention norm 2
@@ -464,6 +465,7 @@ class TensorNameMap:
             "layers.{bid}.ffn_norm",                                        # llama-pth
             "encoder.layer.{bid}.output.LayerNorm",                         # bert
             "language_model.encoder.layers.{bid}.post_attention_layernorm", # persimmon
+            "model.layers.{bid}.ln2",                                       # yi
         ),
 
         # Feed-forward up