From a035de2fdd60d1ce1cce1fb8f4b629f0d68bc258 Mon Sep 17 00:00:00 2001
From: Ettore Di Giacinto <mudler@users.noreply.github.com>
Date: Mon, 15 May 2023 08:15:01 +0200
Subject: [PATCH] tests: add rwkv (#261)

---
 Makefile                            |  2 ++
 api/api_test.go                     | 16 ++++++++++++++--
 tests/fixtures/rwkv.yaml            | 19 +++++++++++++++++++
 tests/fixtures/rwkv_chat.tmpl       | 13 +++++++++++++
 tests/fixtures/rwkv_completion.tmpl |  1 +
 5 files changed, 49 insertions(+), 2 deletions(-)
 create mode 100644 tests/fixtures/rwkv.yaml
 create mode 100644 tests/fixtures/rwkv_chat.tmpl
 create mode 100644 tests/fixtures/rwkv_completion.tmpl

diff --git a/Makefile b/Makefile
index 4cf3a76..683b5d3 100644
--- a/Makefile
+++ b/Makefile
@@ -187,6 +187,8 @@ test-models/testmodel:
 	wget https://huggingface.co/ggerganov/whisper.cpp/resolve/main/ggml-base.en.bin -O test-models/whisper-en
 	wget https://huggingface.co/skeskinen/ggml/resolve/main/all-MiniLM-L6-v2/ggml-model-q4_0.bin -O test-models/bert
 	wget https://cdn.openai.com/whisper/draft-20220913a/micro-machines.wav -O test-dir/audio.wav
+	wget https://huggingface.co/imxcstar/rwkv-4-raven-ggml/resolve/main/RWKV-4-Raven-1B5-v11-Eng99%25-Other1%25-20230425-ctx4096-16_Q4_2.bin -O test-models/rwkv
+	wget https://raw.githubusercontent.com/saharNooby/rwkv.cpp/5eb8f09c146ea8124633ab041d9ea0b1f1db4459/rwkv/20B_tokenizer.json -O test-models/rwkv.tokenizer.json
 	cp tests/fixtures/* test-models
 
 test: prepare test-models/testmodel
diff --git a/api/api_test.go b/api/api_test.go
index def4e20..a05096d 100644
--- a/api/api_test.go
+++ b/api/api_test.go
@@ -47,7 +47,7 @@ var _ = Describe("API test", func() {
 		It("returns the models list", func() {
 			models, err := client.ListModels(context.TODO())
 			Expect(err).ToNot(HaveOccurred())
-			Expect(len(models.Models)).To(Equal(7))
+			Expect(len(models.Models)).To(Equal(10))
 		})
 		It("can generate completions", func() {
 			resp, err := client.CreateCompletion(context.TODO(), openai.CompletionRequest{Model: "testmodel", Prompt: "abcdedfghikl"})
@@ -123,6 +123,18 @@ var _ = Describe("API test", func() {
 			Expect(err).ToNot(HaveOccurred())
 			Expect(resp2.Data[0].Embedding).To(Equal(sunEmbedding))
 		})
+
+		Context("backends", func() {
+			It("runs rwkv", func() {
+				if runtime.GOOS != "linux" {
+					Skip("test supported only on linux")
+				}
+				resp, err := client.CreateCompletion(context.TODO(), openai.CompletionRequest{Model: "rwkv_test", Prompt: "Count up to five: one, two, three, four,"})
+				Expect(err).ToNot(HaveOccurred())
+				Expect(len(resp.Choices) > 0).To(BeTrue())
+				Expect(resp.Choices[0].Text).To(Equal(" five."))
+			})
+		})
 	})
 
 	Context("Config file", func() {
@@ -149,7 +161,7 @@ var _ = Describe("API test", func() {
 
 			models, err := client.ListModels(context.TODO())
 			Expect(err).ToNot(HaveOccurred())
-			Expect(len(models.Models)).To(Equal(9))
+			Expect(len(models.Models)).To(Equal(12))
 		})
 		It("can generate chat completions from config file", func() {
 			resp, err := client.CreateChatCompletion(context.TODO(), openai.ChatCompletionRequest{Model: "list1", Messages: []openai.ChatCompletionMessage{openai.ChatCompletionMessage{Role: "user", Content: "abcdedfghikl"}}})
diff --git a/tests/fixtures/rwkv.yaml b/tests/fixtures/rwkv.yaml
new file mode 100644
index 0000000..d78f5cf
--- /dev/null
+++ b/tests/fixtures/rwkv.yaml
@@ -0,0 +1,19 @@
+name: rwkv_test
+parameters:
+  model: rwkv
+  top_k: 80
+  temperature: 0.9
+  max_tokens: 100
+  top_p: 0.8
+context_size: 1024
+threads: 14
+backend: "rwkv"
+cutwords:
+- "Bob:.*"
+roles:
+  user: "Bob:"
+  system: "Alice:"
+  assistant: "Alice:"
+template:
+  completion: rwkv_completion
+  chat: rwkv_chat
\ No newline at end of file
diff --git a/tests/fixtures/rwkv_chat.tmpl b/tests/fixtures/rwkv_chat.tmpl
new file mode 100644
index 0000000..d2c0511
--- /dev/null
+++ b/tests/fixtures/rwkv_chat.tmpl
@@ -0,0 +1,13 @@
+The following is a verbose detailed conversation between Bob and a woman, Alice. Alice is intelligent, friendly and likeable. Alice is likely to agree with Bob.
+
+Bob: Hello Alice, how are you doing?
+
+Alice: Hi Bob! Thanks, I'm fine. What about you?
+
+Bob: I am very good! It's nice to see you. Would you mind me chatting with you for a while?
+
+Alice: Not at all! I'm listening.
+
+{{.Input}}
+
+Alice: 
\ No newline at end of file
diff --git a/tests/fixtures/rwkv_completion.tmpl b/tests/fixtures/rwkv_completion.tmpl
new file mode 100644
index 0000000..8450377
--- /dev/null
+++ b/tests/fixtures/rwkv_completion.tmpl
@@ -0,0 +1 @@
+Complete the following sentence: {{.Input}} 
\ No newline at end of file