diff --git a/vespa/test_application.py b/vespa/test_application.py index 88fe7153..1f6117a1 100644 --- 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 index 42ea70bd..6a8a9636 100644 --- 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 index 1b2ad028..00536480 100644 --- 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