tomaarsen HF staff commited on
Commit
d901c3c
1 Parent(s): dbc4a34

Add `save_in_root` to the custom Module to allow saving & loading locally

Browse files

Hello!

## 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

Files changed (1) hide show
  1. custom_st.py +3 -0
custom_st.py CHANGED
@@ -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',