apunkt/plano
1
0
Fork 0
mirror of https://github.com/katanemo/plano.git synced 2026-05-03 04:42:49 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions

add test and pre commit hook

Browse source
This commit is contained in:
cotran 2024-10-16 14:15:13 -07:00
parent cf612409f0
commit ee36e403d1
9 changed files with 351 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

9
model_server/app/loader.py
View file

@ -7,6 +7,10 @@ from optimum.onnxruntime import (
ORTModelForSequenceClassification,
)
import app.commons.utilities as utils
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from optimum.intel import OVModelForSequenceClassification
logger = utils.get_model_server_logger()
@ -64,12 +68,9 @@ def get_prompt_guard(model_name):
logger.info("Loading Guard Model...")
if glb.DEVICE == "cpu":
from optimum.intel import OVModelForSequenceClassification
model_class = OVModelForSequenceClassification
elif glb.DEVICE == "gpu":
import torch
from transformers import AutoModelForSequenceClassification
else:
model_class = AutoModelForSequenceClassification