refactor(embeddingref): replace STactor(embedding): replace STEmbeddi…

…ng withEmbedding with LocalEmbedding LocalEmbedding UpdatedUpdated the the Embed EmbeddingEngine to use LocalEmbedding instead of STEmbedding for better consistency and clarity.dingEngine to use LocalEmbedding instead of STEmbedding for better consistency and clarity.
unit-mesh · Jul 6, 2024 · fb267c5 · fb267c5
1 parent 769a258
commit fb267c5
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 5 deletions.
diff --git a/llm-modules/local-embedding/src/main/kotlin/cc/unitmesh/cf/STEmbedding.kt b/llm-modules/local-embedding/src/main/kotlin/cc/unitmesh/cf/STEmbedding.kt
@@ -10,4 +10,10 @@ class STEmbedding(
     private val session: OrtSession,
     private val env: OrtEnvironment,
 ) : LocalEmbedding(tokenizer, session, env) {
+
+    companion object {
+        fun create(): LocalEmbedding {
+            return LocalEmbedding.create()
+        }
+    }
 }
diff --git a/rag-modules/rag-script/src/main/kotlin/cc/unitmesh/rag/EmbeddingEngine.kt b/rag-modules/rag-script/src/main/kotlin/cc/unitmesh/rag/EmbeddingEngine.kt
@@ -1,6 +1,6 @@
 package cc.unitmesh.rag
 
-import cc.unitmesh.cf.STEmbedding
+import cc.unitmesh.cf.LocalEmbedding
 import cc.unitmesh.nlp.embedding.Embedding
 import cc.unitmesh.nlp.embedding.EmbeddingProvider
 import cc.unitmesh.nlp.embedding.text.EnglishTextEmbeddingProvider
@@ -13,14 +13,14 @@ enum class EngineType {
 
 class EmbeddingEngine(private val engine: EngineType = EngineType.SentenceTransformers) {
     var provider: EmbeddingProvider = when (engine) {
-        EngineType.SentenceTransformers -> SentenceTransformersEmbedding()
+        EngineType.SentenceTransformers -> LocalTransformersEmbedding()
         EngineType.EnglishTextEmbedding -> EnglishTextEmbeddingProvider()
         EngineType.TextEmbeddingAda -> TODO()
     }
 }
 
-class SentenceTransformersEmbedding : EmbeddingProvider {
-    private val semantic = STEmbedding.create()
+class LocalTransformersEmbedding : EmbeddingProvider {
+    private val semantic = LocalEmbedding.create()
     override fun embed(texts: List<String>): List<Embedding> {
         return texts.map {
             semantic.embed(it).toList()

diff --git a/server/src/test/kotlin/RagIntegrationTests.kt b/server/src/test/kotlin/RagIntegrationTests.kt
@@ -1,3 +1,4 @@
+import cc.unitmesh.cf.LocalEmbedding
 import cc.unitmesh.cf.STEmbedding
 import cc.unitmesh.cf.infrastructure.llms.embedding.SentenceTransformersEmbedding
 import cc.unitmesh.nlp.embedding.Embedding
@@ -13,7 +14,7 @@ import io.kotest.matchers.shouldBe
 import org.junit.jupiter.api.Test
 
 class RagIntegrationTests {
-    val semantic = STEmbedding.create()
+    val semantic = LocalEmbedding.create()
 
     private val embeddingProvider = SentenceTransformersEmbedding()