Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Update Model methods #71

Merged
merged 1 commit into from
Jan 9, 2024
Merged

Update Model methods #71

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 9 additions & 20 deletions hub_sdk/modules/models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ def is_resumable(self) -> bool:
Returns:
bool: True if resumable, False otherwise.
"""
return self.data.get("hasLastWeights", False)
return self.data.get("status") == "training"

def has_best_weights(self) -> bool:
"""
Expand All @@ -120,7 +120,7 @@ def has_best_weights(self) -> bool:
Returns:
bool: True if best weights available, False otherwise.
"""
return self.data.get("hasBestWeights", False)
return self.is_trained() and bool(self.data.get("weights"))

def is_pretrained(self) -> bool:
"""
Expand All @@ -138,7 +138,7 @@ def is_trained(self) -> bool:
Returns:
bool: True if trained, False otherwise.
"""
return self.data.get("isTrained", False)
return self.data.get("status") == "trained"

def is_custom(self) -> bool:
"""
Expand All @@ -156,8 +156,7 @@ def get_architecture(self) -> str:
Returns:
str or None: The architecture name followed by '.yaml' or None if not available.
"""
name = self.data.get("lineage", {}).get("architecture", {}).get("name")
return f"{name}.yaml" if name else None
return self.data.get("cfg")

def get_dataset_url(self) -> str:
"""
Expand All @@ -166,12 +165,7 @@ def get_dataset_url(self) -> str:
Returns:
str or None: The URL of the dataset or None if not available.
"""
resp = requests.post(
f"{HUB_FUNCTIONS_ROOT}/v1/storage",
json={"collection": "models", "docId": self.id, "object": "dataset"},
headers=self.headers,
)
return resp.json().get("data", {}).get("url")
return self.data.get("data")

def get_weights_url(self, weight: str = "best"):
"""
Expand All @@ -183,15 +177,10 @@ def get_weights_url(self, weight: str = "best"):
Returns:
str or None: The URL of the specified weights or None if not available.
"""
if weight != "parent" or self.is_custom():
resp = requests.post(
f"{HUB_FUNCTIONS_ROOT}/v1/storage",
json={"collection": "models", "docId": self.id, "object": weight},
headers=self.headers,
)
return resp.json().get("data", {}).get("url")
else:
return self.data.get("lineage", {}).get("parent", {}).get("url")
if weight == "last":
return self.data("resume")

return self.data.get("weights")

def delete(self, hard: bool = False) -> dict:
"""
Expand Down