Add `save_in_root` to the custom Module to allow saving & loading locally
Browse filesHello!
## Pull Request overview
* Add `save_in_root` to the custom Module to allow saving & loading locally
## Details
Just like in [jina-embeddings-v3](https://huggingface.co/jinaai/jina-embeddings-v3/blob/main/custom_st.py#L36), we can add `save_in_root = True` to ensure that this module is saved in the root of the directory. This is required because `transformers` expects any custom modeling files (such as `custom_st.py` itself) to be in the root. Without this change, we'll get this after saving with Sentence Transformers:
![image.png](https://cdn-uploads.huggingface.co/production/uploads/6317233cc92fd6fee317e030/FrjDeDQZKbosavf40WXju.png)
With this change in place, we'll instead get this:
![image.png](https://cdn-uploads.huggingface.co/production/uploads/6317233cc92fd6fee317e030/KRANEuLM4kFjDA7KdXIzR.png)
and then you'll be able to save and load this correctly in Sentence Transformers.
- Tom Aarsen
- custom_st.py +3 -0
@@ -12,6 +12,9 @@ from transformers import AutoConfig, AutoImageProcessor, AutoModel, AutoTokenize
|
|
12 |
|
13 |
|
14 |
class Transformer(nn.Module):
|
|
|
|
|
|
|
15 |
def __init__(
|
16 |
self,
|
17 |
model_name_or_path: str = 'jinaai/jina-clip-v2',
|
|
|
12 |
|
13 |
|
14 |
class Transformer(nn.Module):
|
15 |
+
|
16 |
+
save_in_root: bool = True
|
17 |
+
|
18 |
def __init__(
|
19 |
self,
|
20 |
model_name_or_path: str = 'jinaai/jina-clip-v2',
|