apunkt/plano
1
0
Fork 0
mirror of https://github.com/katanemo/plano.git synced 2026-04-29 19:06:34 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

Improve cli (#179)

Browse source
This commit is contained in:
Adil Hafeez 2024-10-10 17:44:41 -07:00 • committed by GitHub
parent ceca0dba28
commit 7d5f760884
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
18 changed files with 611 additions and 445 deletions
Download patch file Download diff file Expand all files Collapse all files

9
model_server/app/loader.py
View file

@ -6,12 +6,15 @@ from optimum.onnxruntime import (
ORTModelForFeatureExtraction,
ORTModelForSequenceClassification,
)
import app.commons.utilities as utils
logger = utils.get_model_server_logger()
def get_embedding_model(
model_name=os.getenv("MODELS", "katanemo/bge-large-en-v1.5"),
):
print("Loading Embedding Model...")
logger.info("Loading Embedding Model...")
if glb.DEVICE != "cuda":
model = ORTModelForFeatureExtraction.from_pretrained(
@ -32,7 +35,7 @@ def get_embedding_model(
def get_zero_shot_model(
model_name=os.getenv("ZERO_SHOT_MODELS", "katanemo/bart-large-mnli"),
):
print("Loading Zero-shot Model...")
logger.info("Loading Zero-shot Model...")
if glb.DEVICE != "cuda":
model = ORTModelForSequenceClassification.from_pretrained(
@ -58,7 +61,7 @@ def get_zero_shot_model(
def get_prompt_guard(model_name, hardware_config="cpu"):
print("Loading Guard Model...")
logger.info("Loading Guard Model...")
if hardware_config == "cpu":
from optimum.intel import OVModelForSequenceClassification