Skip to content

Commit

Permalink
chore: port to Xenova Jina source (#102)
Browse files Browse the repository at this point in the history
* chore: xenova jina

* chore: try recusive model location

* chore: updated doc string, blob pattern
  • Loading branch information
Anush008 authored Jan 30, 2024
1 parent ede507e commit 87decb0
Show file tree
Hide file tree
Showing 2 changed files with 18 additions and 11 deletions.
25 changes: 16 additions & 9 deletions fastembed/embedding.py
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,21 @@ def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:
yield b


def locate_model_file(model_dir: Path, file_names: List[str]):
"""
Find model path for both TransformerJS style `onnx` subdirectory structure and direct model weights structure used by Optimum and Qdrant
"""
if not model_dir.is_dir():
raise ValueError(f"Provided model path '{model_dir}' is not a directory.")

for path in model_dir.rglob("*.onnx"):
for file_name in file_names:
if path.is_file() and path.name == file_name:
return path

raise ValueError(f"Could not find either of {', '.join(file_names)} in {model_dir}")


def normalize(input_array, p=2, dim=1, eps=1e-12):
# Calculate the Lp norm along the specified dimension
norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
Expand Down Expand Up @@ -92,19 +107,11 @@ def __init__(
):
self.path = path
self.model_name = model_name
model_path = self.path / "model.onnx"
optimized_model_path = self.path / "model_optimized.onnx"
model_path = locate_model_file(self.path, ["model.onnx", "model_optimized.onnx"])

# List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
onnx_providers = ["CPUExecutionProvider"]

if not model_path.exists():
# Rename file model_optimized.onnx to model.onnx if it exists
if optimized_model_path.exists():
optimized_model_path.rename(model_path)
else:
raise ValueError(f"Could not find model.onnx in {self.path}")

# Hacky support for multilingual model
self.exclude_token_type_ids = False
if model_name == "intfloat/multilingual-e5-large":
Expand Down
4 changes: 2 additions & 2 deletions fastembed/models.json
Original file line number Diff line number Diff line change
Expand Up @@ -83,7 +83,7 @@
"description": " English embedding model supporting 8192 sequence length",
"size_in_GB": 0.55,
"hf_sources": [
"jinaai/jina-embeddings-v2-base-en"
"xenova/jina-embeddings-v2-base-en"
],
"compressed_url_sources": []
},
Expand All @@ -93,7 +93,7 @@
"description": " English embedding model supporting 8192 sequence length",
"size_in_GB": 0.13,
"hf_sources": [
"jinaai/jina-embeddings-v2-small-en"
"xenova/jina-embeddings-v2-small-en"
],
"compressed_url_sources": []
},
Expand Down

0 comments on commit 87decb0

Please sign in to comment.