fix gpu off error and add GetEmbedding api

Author: Cyberhan123

embed_windows.go

@@ -37,21 +37,23 @@ func getDl(gpu bool) []byte {
 			log.Println(err)
 		}
 		if supportGpuTable[gpuInfo] != nil {
-			return supportGpuTable[gpu]
-		} else {
-			log.Println("GPU not support, use CPU instead.")
+			return supportGpuTable[gpuInfo]
 		}
+		log.Println("GPU not support, use CPU instead.")
 	}
 	if cpu.X86.HasAVX512 {
+		log.Println("Use CPU AVX512 instead.")
 		return libRwkvAvx512
 	}
 	if cpu.X86.HasAVX2 {
+		log.Println("Use CPU AVX2 instead.")
 		return libRwkvAvx2
 	}
 	if cpu.X86.HasAVX {
+		log.Println("Use CPU AVX instead.")
 		return libRwkvAvx
 	}
-	//return libRwkvHipLAS
+
 	panic("Automatic loading of dynamic library failed, please use `NewRwkvModel` method load manually. ")
 	return nil
 }

rwkv.go

@@ -77,7 +77,7 @@ func NewRwkvModel(dylibPath string, options RwkvOptions) (*RwkvModel, error) {
 	}
 
 	if options.GpuEnable {
-		log.Printf("You are about to offload your model to the GPU. " +
+		log.Printf("If you want to try offload your model to the GPU. " +
 			"Please confirm the size of your GPU memory to prevent memory overflow." +
 			"If the model is larger than GPU memory, please specify the layers to offload.")
 	}
@@ -181,6 +181,29 @@ func (s *RwkvState) Predict(input string) (string, error) {
 	return s.generateResponse(nil)
 }
 
+// GetEmbedding give the model embedding.
+// the embedding in rwkv is hidden state the len is n_emb*5*n_layer=46080.
+// So if distillation is true, we split len to n_emb = 768
+func (s *RwkvState) GetEmbedding(input string, distill bool) ([]float32, error) {
+	encode, err := s.rwkvModel.tokenizer.Encode(input)
+
+	for _, token := range encode {
+		err = s.rwkvModel.cRwkv.RwkvEval(s.rwkvModel.ctx, uint32(token), s.state, s.state, s.logits)
+		if err != nil {
+			return nil, err
+		}
+	}
+	// we should keep state clean
+	nState := s.rwkvModel.cRwkv.RwkvGetStateLength(s.rwkvModel.ctx)
+	if distill {
+		nState = s.rwkvModel.cRwkv.RwkvGetNEmbedding(s.rwkvModel.ctx)
+	}
+	emb := make([]float32, nState)
+	copy(emb, s.state)
+	s.state = make([]float32, nState)
+	return emb, nil
+}
+
 func (s *RwkvState) PredictStream(input string, output chan string) {
 	go func() {
 		err := s.handelInput(input)

rwkv_test.go

@@ -277,3 +277,51 @@ func TestChat(t *testing.T) {
 	})
 
 }
+
+func TestRwkvState_GetEmbedding(t *testing.T) {
+	rwkv, err := NewRwkvAutoModel(RwkvOptions{
+		MaxTokens:     500,
+		TokenizerType: Normal, //or World
+		PrintError:    true,
+		CpuThreads:    2,
+		GpuEnable:     true,
+	})
+	defer rwkv.Close()
+
+	err = rwkv.LoadFromFile("./models/RWKV-4b-Pile-171M-20230202-7922-f16.bin")
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	ctx, err := rwkv.InitState()
+	if err != nil {
+		t.Error(err)
+		return
+	}
+	nb := ctx.rwkvModel.cRwkv.RwkvGetNEmbedding(ctx.rwkvModel.ctx)
+	t.Log(nb)
+	nl := ctx.rwkvModel.cRwkv.RwkvGetNLayer(ctx.rwkvModel.ctx)
+
+	t.Run("hidden state", func(t *testing.T) {
+		embedding, err := ctx.GetEmbedding("hello word", false)
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		t.Log(embedding)
+		t.Log(len(embedding))
+		t.Log(nl)
+		assert(t, len(embedding) == int(nb*5*nl))
+	})
+
+	t.Run("distill hidden state ", func(t *testing.T) {
+		embedding, err := ctx.GetEmbedding("hello word", true)
+		if err != nil {
+			t.Error(err)
+			return
+		}
+		t.Log(embedding)
+		assert(t, len(embedding) == int(nb))
+	})
+
+}