Add a new llama_cpp engine (#14436)

JohnSnowLabs · Oct 18, 2024 · b314f62 · b314f62
1 parent 7eadf9e
commit b314f62
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git a/src/main/scala/com/johnsnowlabs/ml/util/ModelEngine.scala b/src/main/scala/com/johnsnowlabs/ml/util/ModelEngine.scala
@@ -33,6 +33,7 @@ final case object ONNX extends ModelEngine {
   val decoderModel = "decoder_model.onnx"
   val decoderWithPastModel = "decoder_with_past_model.onnx"
 }
+
 final case object Openvino extends ModelEngine {
   val name = "openvino"
   val ovModel = "openvino_model"
@@ -41,6 +42,10 @@ final case object Openvino extends ModelEngine {
   val decoderModelWithPast = "openvino_decoder_with_past_model"
 }
 
+final case object LlamaCPP extends ModelEngine {
+  val name = "llama_cpp"
+}
+
 final case object Unknown extends ModelEngine {
   val name = "unk"
 }
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/seq2seq/AutoGGUFModel.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/seq2seq/AutoGGUFModel.scala
@@ -16,9 +16,9 @@
 package com.johnsnowlabs.nlp.annotators.seq2seq
 
 import com.johnsnowlabs.ml.gguf.GGUFWrapper
+import com.johnsnowlabs.ml.util.LlamaCPP
 import com.johnsnowlabs.nlp._
 import com.johnsnowlabs.nlp.util.io.ResourceHelper
-import com.johnsnowlabs.nlp.llama.LlamaModel
 import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.ml.util.Identifiable
 import org.apache.spark.sql.SparkSession
@@ -153,6 +153,8 @@ class AutoGGUFModel(override val uid: String)
     this
   }
 
+  private[johnsnowlabs] def setEngine(engineName: String): this.type = set(engine, engineName)
+
   override def onWrite(path: String, spark: SparkSession): Unit = {
     super.onWrite(path, spark)
     getModelIfNotSet.saveToFile(path)
@@ -261,6 +263,7 @@ trait ReadAutoGGUFModel {
     val annotatorModel = new AutoGGUFModel()
     annotatorModel
       .setModelIfNotSet(spark, GGUFWrapper.read(spark, localPath))
+      .setEngine(LlamaCPP.name)
 
     val metadata = LlamaModel.getMetadataFromFile(localPath)
     if (metadata.nonEmpty) annotatorModel.setMetadata(metadata)