update integration tests

vespa-engine · Oct 3, 2022 · 1a91da0 · 1a91da0
1 parent 4f0971a
commit 1a91da0
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 13 deletions.
diff --git a/vespa/test_application.py b/vespa/test_application.py
@@ -8,7 +8,7 @@
 from vespa.package import ApplicationPackage, ModelServer, Schema, Document
 from vespa.application import Vespa, parse_labeled_data, parse_feed_df
 from vespa.io import VespaQueryResponse
-from vespa.query import QueryModel, OR, RankProfile
+from learntorank.query import QueryModel, OR, Ranking
 
 
 class TestVespa(unittest.TestCase):
@@ -85,7 +85,7 @@ def test_query(self):
         self.assertDictEqual(
             app.query(
                 query="this is a test",
-                query_model=QueryModel(match_phase=OR(), rank_profile=RankProfile()),
+                query_model=QueryModel(match_phase=OR(), ranking=Ranking()),
                 debug_request=True,
                 hits=10,
             ).request_body,
@@ -99,7 +99,7 @@ def test_query(self):
         self.assertDictEqual(
             app.query(
                 query="this is a test",
-                query_model=QueryModel(match_phase=OR(), rank_profile=RankProfile()),
+                query_model=QueryModel(match_phase=OR(), ranking=Ranking()),
                 debug_request=True,
                 hits=10,
                 recall=("id", [1, 5]),
@@ -330,7 +330,7 @@ def test_collect_training_data_point(self):
                 ),
             ]
         )
-        query_model = QueryModel(rank_profile=RankProfile(list_features=True))
+        query_model = QueryModel(ranking=Ranking(list_features=True))
         data = self.app.collect_training_data_point(
             query="this is a query",
             query_id="123",
@@ -399,7 +399,7 @@ def test_collect_training_data_point_absent_field(self):
                 ),
             ]
         )
-        query_model = QueryModel(rank_profile=RankProfile(list_features=True))
+        query_model = QueryModel(ranking=Ranking(list_features=True))
         data = self.app.collect_training_data_point(
             query="this is a query",
             query_id="123",
@@ -480,7 +480,7 @@ def test_collect_training_data_point_0_recall_hits(self):
                 ),
             ]
         )
-        query_model = QueryModel(rank_profile=RankProfile(list_features=True))
+        query_model = QueryModel(ranking=Ranking(list_features=True))
         data = self.app.collect_training_data_point(
             query="this is a query",
             query_id="123",
@@ -539,7 +539,7 @@ def test_collect_training_data(self):
                 "relevant_docs": [{"id": "abc", "score": 1}],
             }
         ]
-        query_model = QueryModel(rank_profile=RankProfile(list_features=True))
+        query_model = QueryModel(ranking=Ranking(list_features=True))
         data = self.app.collect_training_data(
             labeled_data=labeled_data,
             id_field="vespa_id_field",

diff --git a/vespa/test_integration_docker.py b/vespa/test_integration_docker.py
@@ -18,9 +18,9 @@
 )
 from vespa.deployment import VespaDocker
 from vespa.ml import BertModelConfig, SequenceClassification
-from vespa.query import QueryModel, RankProfile as Ranking, OR, QueryRankingFeature
 from vespa.gallery import QuestionAnswering, TextSearch
 from vespa.application import VespaSync
+from learntorank.query import QueryModel, Ranking, OR, QueryRankingFeature
 
 CONTAINER_STOP_TIMEOUT = 600
 
@@ -940,7 +940,7 @@ def bert_model_input_and_output(
                     )
                 ],
                 match_phase=OR(),
-                rank_profile=Ranking(name="pretrained_bert_tiny"),
+                ranking=Ranking(name="pretrained_bert_tiny"),
             ),
         )
         vespa_input_ids = self._parse_vespa_tensor(result.hits[0], "input_ids")
@@ -1056,7 +1056,7 @@ def setUp(self) -> None:
         ]
         self.query_batch = ["Give me title 1", "Give me title 2"]
         self.query_model = QueryModel(
-            match_phase=OR(), rank_profile=Ranking(name="default", list_features=False)
+            match_phase=OR(), ranking=Ranking(name="default", list_features=False)
         )
         self.queries_first_hit = ["this is title 1", "this is title 2"]
 
@@ -1190,7 +1190,7 @@ def test_check_duplicated_features(self):
             relevant_id="1",
             id_field="id",
             query_model=QueryModel(
-                match_phase=OR(), rank_profile=Ranking(name="bm25", list_features=True)
+                match_phase=OR(), ranking=Ranking(name="bm25", list_features=True)
             ),
             number_additional_docs=10,
             fields=["rankfeatures", "summaryfeatures"],
@@ -1230,7 +1230,7 @@ def test_store_vespa_features(self):
             labeled_data=labeled_data,
             id_field="id",
             query_model=QueryModel(
-                match_phase=OR(), rank_profile=Ranking(name="bm25", list_features=True)
+                match_phase=OR(), ranking=Ranking(name="bm25", list_features=True)
             ),
             number_additional_docs=2,
             fields=["rankfeatures", "summaryfeatures"],

diff --git a/vespa/test_integration_vespa_cloud.py b/vespa/test_integration_vespa_cloud.py
@@ -6,7 +6,6 @@
 from pandas import DataFrame
 from cryptography.hazmat.primitives import serialization
 from vespa.application import Vespa
-from vespa.query import QueryModel, OR
 from vespa.gallery import TextSearch
 from vespa.deployment import VespaCloud
 from vespa.test_integration_docker import (
@@ -16,6 +15,7 @@
     create_qa_application_package,
     create_sequence_classification_task,
 )
+from learntorank.query import QueryModel, OR
 
 CFG_SERVER_START_TIMEOUT = 300
 APP_INIT_TIMEOUT = 300