philschmid
/

all-MiniLM-L6-v2-optimum-embeddings

sentence-embeddings

endpoints-template

Inference Endpoints

Model card Files Files and versions Community

philschmid HF staff commited on Sep 16, 2022

Commit

bf24a3c

·

1 Parent(s): 3fcf89c

Update pipeline.py

Files changed (1) hide show

pipeline.py +1 -1

pipeline.py CHANGED Viewed

@@ -11,7 +11,7 @@ def mean_pooling(model_output, attention_mask):
     return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
-class PreTrainedPipeline():
     def __init__(self, path=""):
         # load the optimized model
         self.model = ORTModelForFeatureExtraction.from_pretrained(path, file_name="model-quantized.onnx")

     return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)
+class EndpointHandler():
     def __init__(self, path=""):
         # load the optimized model
         self.model = ORTModelForFeatureExtraction.from_pretrained(path, file_name="model-quantized.onnx")