From 87ea2d96fde3540d1a55aa1052d86046a509b969 Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Tue, 26 Sep 2023 17:43:39 -0700
Subject: [PATCH] Add a note about RWKV loader

---
 modules/RWKV.py   |  5 +++++
 modules/models.py | 20 ++++++++++++--------
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/modules/RWKV.py b/modules/RWKV.py
index 35d69986..39487c66 100644
--- a/modules/RWKV.py
+++ b/modules/RWKV.py
@@ -1,3 +1,8 @@
+'''
+This loader is not currently maintained as RWKV can now be loaded
+through the transformers library.
+'''
+
 import copy
 import os
 from pathlib import Path
diff --git a/modules/models.py b/modules/models.py
index 06dfe994..253c998b 100644
--- a/modules/models.py
+++ b/modules/models.py
@@ -211,14 +211,6 @@ def huggingface_loader(model_name):
     return model
 
 
-def RWKV_loader(model_name):
-    from modules.RWKV import RWKVModel, RWKVTokenizer
-
-    model = RWKVModel.from_pretrained(Path(f'{shared.args.model_dir}/{model_name}'), dtype="fp32" if shared.args.cpu else "bf16" if shared.args.bf16 else "fp16", device="cpu" if shared.args.cpu else "cuda")
-    tokenizer = RWKVTokenizer.from_pretrained(Path(shared.args.model_dir))
-    return model, tokenizer
-
-
 def llamacpp_loader(model_name):
     from modules.llamacpp_model import LlamaCppModel
 
@@ -335,6 +327,18 @@ def ExLlamav2_HF_loader(model_name):
     return Exllamav2HF.from_pretrained(model_name)
 
 
+def RWKV_loader(model_name):
+    '''
+    This loader is not currently maintained as RWKV can now be loaded
+    through the transformers library.
+    '''
+    from modules.RWKV import RWKVModel, RWKVTokenizer
+
+    model = RWKVModel.from_pretrained(Path(f'{shared.args.model_dir}/{model_name}'), dtype="fp32" if shared.args.cpu else "bf16" if shared.args.bf16 else "fp16", device="cpu" if shared.args.cpu else "cuda")
+    tokenizer = RWKVTokenizer.from_pretrained(Path(shared.args.model_dir))
+    return model, tokenizer
+
+
 def get_max_memory_dict():
     max_memory = {}
     if shared.args.gpu_memory: