Change pretrained AutoGGUFReranking model

DevinTDHa · DevinTDHa · commit d20e35933d36 · 2025-09-01T12:34:10.000+02:00
diff --git a/docs/en/annotator_entries/AutoGGUFReranker.md b/docs/en/annotator_entries/AutoGGUFReranker.md
@@ -33,7 +33,7 @@ val reranker = AutoGGUFReranker.pretrained()
   .setQuery("A man is eating pasta.")
 ```
 
-The default model is `"bge-reranker-v2-m3-Q4_K_M"`, if no name is provided.
+The default model is `"bge_reranker_v2_m3_Q4_K_M"`, if no name is provided.
 
 For available pretrained models please see the [Models Hub](https://sparknlp.org/models).
 
@@ -105,7 +105,7 @@ val document = new DocumentAssembler()
   .setOutputCol("document")
 
 val reranker = AutoGGUFReranker
-  .pretrained("bge-reranker-v2-m3-Q4_K_M")
+  .pretrained()
   .setInputCols("document")
   .setOutputCol("reranked_documents")
   .setBatchSize(4)
diff --git a/docs/en/annotator_entries/GGUFRankingFinisher.md b/docs/en/annotator_entries/GGUFRankingFinisher.md
@@ -85,7 +85,7 @@ val documentAssembler = new DocumentAssembler()
 
 // Reranker
 val reranker = AutoGGUFReranker
-  .pretrained("bge-reranker-v2-m3-Q4_K_M")
+  .pretrained()
   .setInputCols("document")
   .setOutputCol("reranked_documents")
   .setQuery("A man is eating pasta.")
diff --git a/examples/python/llama.cpp/GGUFRankingFinisher_for_AutoGGUFReranker.ipynb b/examples/python/llama.cpp/GGUFRankingFinisher_for_AutoGGUFReranker.ipynb
@@ -136,10 +136,7 @@
     "document_assembler = DocumentAssembler().setInputCol(\"text\").setOutputCol(\"document\")\n",
     "\n",
     "auto_gguf_model = (\n",
-    "    AutoGGUFReranker.loadSavedModel(\n",
-    "        \"/home/ducha/Workspace/scala/spark-nlp-release/tmp_autogguf_reranker/bge-reranker-v2-m3-q4_k_m.gguf\",\n",
-    "        spark,\n",
-    "    )\n",
+    "    AutoGGUFReranker.pretrained()\n",
     "    .setInputCols(\"document\")\n",
     "    .setOutputCol(\"reranked_documents\")\n",
     "    .setQuery(\"A man is eating pasta.\")\n",
diff --git a/python/sparknlp/annotator/seq2seq/auto_gguf_reranker.py b/python/sparknlp/annotator/seq2seq/auto_gguf_reranker.py
@@ -47,7 +47,7 @@ class AutoGGUFReranker(AnnotatorModel, HasBatchedAnnotate, HasLlamaCppProperties
     ...     .setOutputCol("reranked_documents") \\
     ...     .setQuery("A man is eating pasta.")
 
-    The default model is ``"bge-reranker-v2-m3-Q4_K_M"``, if no name is provided.
+    The default model is ``"bge_reranker_v2_m3_Q4_K_M"``, if no name is provided.
 
     For extended examples of usage, see the
     `AutoGGUFRerankerTest <https://github.com/JohnSnowLabs/spark-nlp/tree/master/src/test/scala/com/johnsnowlabs/nlp/annotators/seq2seq/AutoGGUFRerankerTest.scala>`__
@@ -222,7 +222,7 @@ class AutoGGUFReranker(AnnotatorModel, HasBatchedAnnotate, HasLlamaCppProperties
     >>> document = DocumentAssembler() \\
     ...     .setInputCol("text") \\
     ...     .setOutputCol("document")
-    >>> reranker = AutoGGUFReranker.pretrained("bge-reranker-v2-m3-Q4_K_M") \\
+    >>> reranker = AutoGGUFReranker.pretrained() \\
     ...     .setInputCols(["document"]) \\
     ...     .setOutputCol("reranked_documents") \\
     ...     .setBatchSize(4) \\
@@ -307,13 +307,13 @@ def loadSavedModel(folder, spark_session):
         return AutoGGUFReranker(java_model=jModel)
 
     @staticmethod
-    def pretrained(name="bge-reranker-v2-m3-Q4_K_M", lang="en", remote_loc=None):
+    def pretrained(name="bge_reranker_v2_m3_Q4_K_M", lang="en", remote_loc=None):
         """Downloads and loads a pretrained model.
 
         Parameters
         ----------
         name : str, optional
-            Name of the pretrained model, by default "bge-reranker-v2-m3-Q4_K_M"
+            Name of the pretrained model, by default "bge_reranker_v2_m3_Q4_K_M"
         lang : str, optional
             Language of the pretrained model, by default "en"
         remote_loc : str, optional
diff --git a/python/sparknlp/base/gguf_ranking_finisher.py b/python/sparknlp/base/gguf_ranking_finisher.py
@@ -65,7 +65,7 @@ class GGUFRankingFinisher(AnnotatorTransformer):
     >>> documentAssembler = DocumentAssembler() \\
     ...     .setInputCol("text") \\
     ...     .setOutputCol("document")
-    >>> reranker = AutoGGUFReranker.pretrained("bge-reranker-v2-m3-Q4_K_M") \\
+    >>> reranker = AutoGGUFReranker.pretrained() \\
     ...     .setInputCols("document") \\
     ...     .setOutputCol("reranked_documents") \\
     ...     .setQuery("A man is eating pasta.")
diff --git a/python/test/annotator/seq2seq/auto_gguf_reranker_test.py b/python/test/annotator/seq2seq/auto_gguf_reranker_test.py
@@ -47,9 +47,10 @@ def runTest(self):
 
         # Use a local model path for testing - in real scenarios, use pretrained()
         model_path = "/tmp/bge-reranker-v2-m3-Q4_K_M.gguf"
-        
+
         # Skip test if model file doesn't exist
         import os
+
         if not os.path.exists(model_path):
             self.skipTest(f"Model file not found: {model_path}")
 
@@ -104,33 +105,29 @@ def runTest(self):
             DocumentAssembler().setInputCol("text").setOutputCol("document")
         )
 
-        # Test with pretrained model (may not be available in test environment)
-        try:
-            reranker = (
-                AutoGGUFReranker.pretrained("bge-reranker-v2-m3-Q4_K_M")
-                .setInputCols("document")
-                .setOutputCol("reranked_documents")
-                .setBatchSize(2)
-                .setQuery(self.query)
-            )
+        reranker = (
+            AutoGGUFReranker.pretrained()
+            .setInputCols("document")
+            .setOutputCol("reranked_documents")
+            .setBatchSize(2)
+            .setQuery(self.query)
+        )
 
-            pipeline = Pipeline().setStages([document_assembler, reranker])
-            results = pipeline.fit(self.data).transform(self.data)
+        pipeline = Pipeline().setStages([document_assembler, reranker])
+        results = pipeline.fit(self.data).transform(self.data)
 
-            # Verify results contain relevance scores
-            collected_results = results.collect()
-            for row in collected_results:
-                annotations = row["reranked_documents"]
-                for annotation in annotations:
-                    self.assertIn("relevance_score", annotation.metadata)
-                    # Relevance score should be a valid number
-                    score = float(annotation.metadata["relevance_score"])
-                    self.assertIsInstance(score, float)
+        # Verify results contain relevance scores
+        collected_results = results.collect()
+        for row in collected_results:
+            annotations = row["reranked_documents"]
+            for annotation in annotations:
+                self.assertIn("relevance_score", annotation.metadata)
+                # Relevance score should be a valid number
+                score = float(annotation.metadata["relevance_score"])
+                self.assertIsInstance(score, float)
+
+        results.show()
 
-            results.show()
-        except Exception as e:
-            # Skip if pretrained model is not available
-            self.skipTest(f"Pretrained model not available: {str(e)}")
 
 @pytest.mark.slow
 class AutoGGUFRerankerMetadataTestSpec(unittest.TestCase):
@@ -139,9 +136,10 @@ def setUp(self):
 
     def runTest(self):
         model_path = "/tmp/bge-reranker-v2-m3-Q4_K_M.gguf"
-        
+
         # Skip test if model file doesn't exist
         import os
+
         if not os.path.exists(model_path):
             self.skipTest(f"Model file not found: {model_path}")
 
@@ -150,10 +148,11 @@ def runTest(self):
         metadata = reranker.getMetadata()
         self.assertIsNotNone(metadata)
         self.assertGreater(len(metadata), 0)
-        
+
         print("Model metadata:")
         print(eval(metadata))
 
+
 #
 # @pytest.mark.slow
 # class AutoGGUFRerankerSerializationTestSpec(unittest.TestCase):
@@ -215,7 +214,7 @@ def runTest(self):
 #         results.select("reranked_documents").show(truncate=False)
 
 
-@pytest.mark.slow 
+@pytest.mark.slow
 class AutoGGUFRerankerErrorHandlingTestSpec(unittest.TestCase):
     def setUp(self):
         self.spark = SparkContextForTest.spark
@@ -229,9 +228,10 @@ def runTest(self):
         data = self.spark.createDataFrame([["Test document"]]).toDF("text")
 
         model_path = "/tmp/bge-reranker-v2-m3-Q4_K_M.gguf"
-        
+
         # Skip test if model file doesn't exist
         import os
+
         if not os.path.exists(model_path):
             self.skipTest(f"Model file not found: {model_path}")
 
@@ -244,7 +244,7 @@ def runTest(self):
         )
 
         pipeline = Pipeline().setStages([document_assembler, reranker])
-        
+
         # This should still work with empty query (based on implementation)
         try:
             results = pipeline.fit(data).transform(data)
@@ -279,9 +279,10 @@ def runTest(self):
         )
 
         model_path = "/tmp/bge-reranker-v2-m3-Q4_K_M.gguf"
-        
+
         # Skip test if model file doesn't exist
         import os
+
         if not os.path.exists(model_path):
             self.skipTest(f"Model file not found: {model_path}")
 
@@ -322,11 +323,11 @@ def runTest(self):
                 self.assertIn("rank", annotation.metadata)
                 self.assertIn("query", annotation.metadata)
                 self.assertEqual(annotation.metadata["query"], self.query)
-                
+
                 # Check that relevance score is normalized (due to minMaxScaling)
                 score = float(annotation.metadata["relevance_score"])
                 self.assertTrue(0.0 <= score <= 1.0)
-                
+
                 # Check that rank is a valid integer
                 rank = int(annotation.metadata["rank"])
                 self.assertIsInstance(rank, int)
@@ -338,7 +339,9 @@ def runTest(self):
             ranks = [int(annotation.metadata["rank"]) for annotation in annotations]
             self.assertEqual(ranks, sorted(ranks))
 
-        print("Pipeline with AutoGGUFReranker and GGUFRankingFinisher completed successfully")
+        print(
+            "Pipeline with AutoGGUFReranker and GGUFRankingFinisher completed successfully"
+        )
         results.select("ranked_documents").show(truncate=False)
 
 
@@ -368,9 +371,10 @@ def runTest(self):
         )
 
         model_path = "/tmp/bge-reranker-v2-m3-Q4_K_M.gguf"
-        
+
         # Skip test if model file doesn't exist
         import os
+
         if not os.path.exists(model_path):
             self.skipTest(f"Model file not found: {model_path}")
 
@@ -396,7 +400,7 @@ def runTest(self):
         results = pipeline.fit(self.data).transform(self.data)
 
         collected_results = results.collect()
-        
+
         # Should have at most 2 results due to topK
         self.assertLessEqual(len(collected_results), 2)
 
@@ -407,7 +411,7 @@ def runTest(self):
                 # Check normalized scores are >= 0.1 threshold
                 score = float(annotation.metadata["relevance_score"])
                 self.assertTrue(0.1 <= score <= 1.0)
-                
+
                 # Check rank metadata exists
                 self.assertIn("rank", annotation.metadata)
                 rank = int(annotation.metadata["rank"])
diff --git a/python/test/base/gguf_ranking_finisher_test.py b/python/test/base/gguf_ranking_finisher_test.py
diff --git a/src/main/scala/com/johnsnowlabs/nlp/annotators/seq2seq/AutoGGUFReranker.scala b/src/main/scala/com/johnsnowlabs/nlp/annotators/seq2seq/AutoGGUFReranker.scala
diff --git a/src/test/scala/com/johnsnowlabs/nlp/annotators/seq2seq/AutoGGUFRerankerTest.scala b/src/test/scala/com/johnsnowlabs/nlp/annotators/seq2seq/AutoGGUFRerankerTest.scala
