YaYaB
/

yb_test_inference_endpoint_text_cls

Inference Endpoints

Model card Files Files and versions Community

YassineB commited on Oct 17, 2022

Commit

5015587

·

1 Parent(s): 423dcb9

Fix requirements

Files changed (1) hide show

handler.py +3 -3

handler.py CHANGED Viewed

@@ -6,13 +6,13 @@ from transformers import pipeline, AutoTokenizer
 class EndpointHandler():
     def __init__(self, path=""):
         # load the optimized model
-        path = "optimum/roberta-base-squad2"
-        model = ORTModelForSequenceClassification.from_pretrained(path)
         tokenizer = AutoTokenizer.from_pretrained(path)
         # create inference pipeline
         self.pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         """
         Args:

 class EndpointHandler():
     def __init__(self, path=""):
         # load the optimized model
+        path = "distilbert-base-uncased-finetuned-sst-2-english"
+        model = ORTModelForSequenceClassification.from_pretrained(path,from_transformers=True)
         tokenizer = AutoTokenizer.from_pretrained(path)
         # create inference pipeline
         self.pipeline = pipeline("text-classification", model=model, tokenizer=tokenizer)
     def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
         """
         Args: