jpohhhh commited on
Commit
dbdba21
·
1 Parent(s): 2758b9b

Try using path...

Browse files
Files changed (1) hide show
  1. handler.py +3 -1
handler.py CHANGED
@@ -2,6 +2,8 @@ from typing import Dict, List, Any
2
  from transformers import AutoTokenizer, AutoModel
3
  from optimum.pipelines import pipeline
4
  from optimum.onnxruntime import ORTModelForFeatureExtraction
 
 
5
 
6
  import torch
7
 
@@ -16,7 +18,7 @@ class EndpointHandler():
16
  # self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
17
  task = "feature-extraction"
18
  tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/msmarco-MiniLM-L-6-v3')
19
- model_regular = ORTModelForFeatureExtraction.from_pretrained("onnx", file_name="model.onnx", from_transformers=False)
20
  self.onnx_extractor = pipeline(task, model=model_regular, tokenizer=tokenizer)
21
  # self.model.to(self.device)
22
  # print("model will run on ", self.device)
 
2
  from transformers import AutoTokenizer, AutoModel
3
  from optimum.pipelines import pipeline
4
  from optimum.onnxruntime import ORTModelForFeatureExtraction
5
+ from pathlib import Path
6
+
7
 
8
  import torch
9
 
 
18
  # self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
19
  task = "feature-extraction"
20
  tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/msmarco-MiniLM-L-6-v3')
21
+ model_regular = ORTModelForFeatureExtraction.from_pretrained(Path("onnx"), file_name="model.onnx", from_transformers=False)
22
  self.onnx_extractor = pipeline(task, model=model_regular, tokenizer=tokenizer)
23
  # self.model.to(self.device)
24
  # print("model will run on ", self.device)