refactor(embedding): modify LocalEmbedding creation method

Modify the LocalEmbedding companion object to accept a ClassLoader parameter for flexibility in model loading. Remove unnecessary comments and instantiation in test class.
unit-mesh · Jul 7, 2024 · f3396e1 · f3396e1
1 parent 84183dd
commit f3396e1
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 10 deletions.
diff --git a/llm-modules/local-embedding/src/main/kotlin/cc/unitmesh/cf/LocalEmbedding.kt b/llm-modules/local-embedding/src/main/kotlin/cc/unitmesh/cf/LocalEmbedding.kt
@@ -67,14 +67,8 @@ open class LocalEmbedding(
 
 
  companion object {
- /**
- * Create a new instance of [LocalEmbedding] with default model.
- * We use official model: [all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2)
- * We can use [optimum](https://github.com/huggingface/optimum) to transform the model to onnx.
- */
- fun create(): LocalEmbedding {
- val classLoader = Thread.currentThread().getContextClassLoader()
 
+ fun create(classLoader: ClassLoader = Thread.currentThread().getContextClassLoader()): LocalEmbedding {
  val tokenizer = loadTokenizer(classLoader)!!
 
  val ortEnv = OrtEnvironment.getEnvironment()

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