Skip to content

Commit cc0a31c

Browse files
authored
[FSTORE-1260] Hsfs should not use hsml client to get embedding model (#1233)
* use hsfs client * fix style
1 parent 725a21e commit cc0a31c

File tree

1 file changed

+15
-3
lines changed

1 file changed

+15
-3
lines changed

python/hsfs/embedding.py

Lines changed: 15 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
#
1616

1717
from dataclasses import dataclass
18+
from hsfs import client
1819
from hsfs.client.exceptions import FeatureStoreException
1920
from typing import Type
2021
import json
@@ -59,16 +60,27 @@ def from_model(cls, model):
5960
# should get from backend because of authorisation check (unshared project etc)
6061
def get_model(self):
6162
try:
62-
from hsml.core.model_api import ModelApi
63+
from hsml.model import Model
6364
except ModuleNotFoundError:
6465
raise FeatureStoreException(
6566
"Model is attached to embedding feature but hsml library is not installed."
6667
"Install hsml library before getting the feature group."
6768
)
6869

69-
return ModelApi().get(
70-
self.model_name, self.model_version, self.model_registry_id
70+
path_params = [
71+
"project",
72+
client.get_instance()._project_id,
73+
"modelregistries",
74+
self.model_registry_id,
75+
"models",
76+
self.model_name + "_" + str(self.model_version),
77+
]
78+
query_params = {"expand": "trainingdatasets"}
79+
80+
model_json = client.get_instance()._send_request(
81+
"GET", path_params, query_params
7182
)
83+
return Model.from_response_json(model_json)
7284

7385
def json(self):
7486
return json.dumps(self, cls=util.FeatureStoreEncoder)

0 commit comments

Comments
 (0)