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