@@ -35,6 +35,15 @@ def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:
35
35
yield b
36
36
37
37
38
+ def locate_model_file (path : Path , file_names : list ):
39
+ for root , dirs , files in path .rglob ("*" ):
40
+ for file_name in file_names :
41
+ if file_name in files :
42
+ return Path (root ) / file_name
43
+
44
+ raise ValueError (f"Could not find model file in { path } " )
45
+
46
+
38
47
def normalize (input_array , p = 2 , dim = 1 , eps = 1e-12 ):
39
48
# Calculate the Lp norm along the specified dimension
40
49
norm = np .linalg .norm (input_array , ord = p , axis = dim , keepdims = True )
@@ -92,32 +101,11 @@ def __init__(
92
101
):
93
102
self .path = path
94
103
self .model_name = model_name
95
- model_path = self .path / "model.onnx"
96
- optimized_model_path = self .path / "model_optimized.onnx"
97
-
98
- xenova_model_path = self .path / "onnx" / "model.onnx"
99
- xenova_optimized_model_path = self .path / "onnx" / "model_optimized.onnx"
104
+ model_path = locate_model_file (self .path , ["model.onnx" , "model_optimized.onnx" ])
100
105
101
106
# List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
102
107
onnx_providers = ["CPUExecutionProvider" ]
103
108
104
- if not model_path .exists ():
105
- # Rename file model_optimized.onnx to model.onnx if it exists
106
- if optimized_model_path .exists ():
107
- optimized_model_path .rename (model_path )
108
-
109
- # Patch for inconsistent repo structure at
110
- # - https://huggingface.co/Xenova/jina-embeddings-v2-small-en
111
- # - https://huggingface.co/Xenova/jina-embeddings-v2-base-en
112
- elif xenova_model_path .exists ():
113
- model_path = xenova_model_path
114
-
115
- elif xenova_optimized_model_path .exists ():
116
- model_path = xenova_optimized_model_path
117
-
118
- else :
119
- raise ValueError (f"Could not find model.onnx in { self .path } " )
120
-
121
109
# Hacky support for multilingual model
122
110
self .exclude_token_type_ids = False
123
111
if model_name == "intfloat/multilingual-e5-large" :
0 commit comments