fix: move to new model definition

BIGmama-technology · Mar 9, 2024 · 5b2f5c6 · 5b2f5c6
1 parent 8aec796
commit 5b2f5c6
Show file tree

Hide file tree

Showing 10 changed files with 33 additions and 18 deletions.
diff --git a/hyko_toolkit/models/multimodal/image_to_text/main.py b/hyko_toolkit/models/multimodal/image_to_text/main.py
@@ -1,6 +1,7 @@
-from metadata import Inputs, Outputs, Params, StartupParams, func
 from transformers import pipeline
 
+from .metadata import Inputs, Outputs, Params, StartupParams, func
+
 
 @func.on_startup
 async def load(startup_params: StartupParams):

diff --git a/hyko_toolkit/models/multimodal/image_to_text/metadata.py b/hyko_toolkit/models/multimodal/image_to_text/metadata.py
@@ -1,10 +1,12 @@
 from pydantic import Field
 
-from hyko_sdk.function import SDKFunction
+from hyko_sdk.definitions import ToolkitModel
 from hyko_sdk.io import Image
 from hyko_sdk.metadata import CoreModel
 
-func = SDKFunction(
+func = ToolkitModel(
+    name="image_to_text",
+    task="multimodal",
     description="Hugging Face Image-To-Text Task",
 )
 

diff --git a/hyko_toolkit/models/multimodal/text_to_image/main.py b/hyko_toolkit/models/multimodal/text_to_image/main.py
@@ -1,8 +1,9 @@
 from diffusers.pipelines.pipeline_utils import DiffusionPipeline
-from metadata import Inputs, Outputs, Params, StartupParams, func
 
 from hyko_sdk.io import Image
 
+from .metadata import Inputs, Outputs, Params, StartupParams, func
+
 generator = None
 
 

diff --git a/hyko_toolkit/models/multimodal/text_to_image/metadata.py b/hyko_toolkit/models/multimodal/text_to_image/metadata.py
@@ -1,10 +1,12 @@
 from pydantic import Field
 
-from hyko_sdk.function import SDKFunction
+from hyko_sdk.definitions import ToolkitModel
 from hyko_sdk.io import Image
 from hyko_sdk.metadata import CoreModel
 
-func = SDKFunction(
+func = ToolkitModel(
+    name="text_to_image",
+    task="multimodal",
     description="Hugging Face Text to Image Task",
 )
 
@@ -25,7 +27,7 @@ class Inputs(CoreModel):
 @func.set_param
 class Params(CoreModel):
     num_inference_steps: int = Field(
-        default=15, description="The number of denoising steps."
+        default=15, description="The number of de-noising steps."
     )
     strength: float = Field(
         default=0.8,

diff --git a/hyko_toolkit/models/multimodal/visual_question_answering/main.py b/hyko_toolkit/models/multimodal/visual_question_answering/main.py
@@ -1,6 +1,7 @@
-from metadata import Inputs, Outputs, Params, StartupParams, func
 from transformers import pipeline
 
+from .metadata import Inputs, Outputs, Params, StartupParams, func
+
 
 @func.on_startup
 async def load(startup_params: StartupParams):

diff --git a/hyko_toolkit/models/multimodal/visual_question_answering/metadata.py b/hyko_toolkit/models/multimodal/visual_question_answering/metadata.py
@@ -1,10 +1,12 @@
 from pydantic import Field
 
-from hyko_sdk.function import SDKFunction
+from hyko_sdk.definitions import ToolkitModel
 from hyko_sdk.io import Image
 from hyko_sdk.metadata import CoreModel
 
-func = SDKFunction(
+func = ToolkitModel(
+    name="visual_question_answering",
+    task="multimodal",
     description="Hugging Face Image-To-Text Task",
 )
 
@@ -29,4 +31,4 @@ class Params(CoreModel):
 @func.set_output
 class Outputs(CoreModel):
     answer: list[str] = Field(..., description="Generated answer")
-    score: list[float] = Field(..., description="Confidance score")
+    score: list[float] = Field(..., description="Confidence score")
diff --git a/hyko_toolkit/models/regression/linear_regression/main.py b/hyko_toolkit/models/regression/linear_regression/main.py
@@ -1,8 +1,9 @@
 import numpy as np
 from fastapi.exceptions import HTTPException
-from metadata import Inputs, Outputs, Params, func
 from sklearn.linear_model import LinearRegression
 
+from .metadata import Inputs, Outputs, Params, func
+
 linear_regression = LinearRegression()
 
 

diff --git a/hyko_toolkit/models/regression/linear_regression/metadata.py b/hyko_toolkit/models/regression/linear_regression/metadata.py
@@ -1,9 +1,11 @@
 from pydantic import Field
 
-from hyko_sdk.function import SDKFunction
+from hyko_sdk.definitions import ToolkitModel
 from hyko_sdk.metadata import CoreModel
 
-func = SDKFunction(
+func = ToolkitModel(
+    name="linear_regression",
+    task="regression",
     description="Predicts a future value based on historical data",
 )
 

diff --git a/hyko_toolkit/models/similarity_search/E5_faiss/main.py b/hyko_toolkit/models/similarity_search/E5_faiss/main.py
@@ -1,7 +1,8 @@
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 from langchain_core.documents.base import Document
-from metadata import Inputs, Outputs, Params, StartupParams, func
+
+from .metadata import Inputs, Outputs, Params, StartupParams, func
 
 
 @func.on_startup

diff --git a/hyko_toolkit/models/similarity_search/E5_faiss/metadata.py b/hyko_toolkit/models/similarity_search/E5_faiss/metadata.py
@@ -1,10 +1,12 @@
 from pydantic import Field
 
-from hyko_sdk.function import SDKFunction
+from hyko_sdk.definitions import ToolkitModel
 from hyko_sdk.metadata import CoreModel
 
-func = SDKFunction(
-    description="Tool for computing similarity scores based on a given threshold."
+func = ToolkitModel(
+    name="E5_faiss",
+    task="similarity_search",
+    description="Tool for computing similarity scores based on a given threshold.",
 )