nbroad
/

deberta-v3-base-prompt-injection-v2-O2

Inference Endpoints

Model card Files Files and versions Community

nbroad HF staff commited on Oct 9

Commit

e8db095

•

1 Parent(s): 9a5e8b9

Update handler.py

Files changed (1) hide show

handler.py +5 -1

handler.py CHANGED Viewed

@@ -7,7 +7,11 @@ from optimum.pipelines import pipeline
 class EndpointHandler():
     def __init__(self, path=""):
         # load the optimized model
-        model = ORTModelForSequenceClassification.from_pretrained(path)
         tokenizer = AutoTokenizer.from_pretrained(path)
         # create inference pipeline
         self.pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer, device=0)

 class EndpointHandler():
     def __init__(self, path=""):
         # load the optimized model
+        model = ORTModelForSequenceClassification.from_pretrained(
+            path,
+            export=False,
+            provider="CUDAExecutionProvider",
+        )
         tokenizer = AutoTokenizer.from_pretrained(path)
         # create inference pipeline
         self.pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer, device=0)