vam commited on
Commit
c0af8d7
1 Parent(s): 990d248

Upload embedding.py

Browse files
Files changed (1) hide show
  1. embedding.py +20 -0
embedding.py ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from sentence_transformers import SentenceTransformer
2
+ from preprocess import meme_attribute, meme_filename, meme_list
3
+ # This model supports two prompts: "s2p_query" and "s2s_query" for sentence-to-passage and sentence-to-sentence tasks, respectively.
4
+ # They are defined in `config_sentence_transformers.json`
5
+
6
+
7
+ # you can also use this model without the features of `use_memory_efficient_attention` and `unpad_inputs`. It can be worked in CPU.
8
+ model = SentenceTransformer(
9
+ "dunzhang/stella_en_400M_v5",
10
+ trust_remote_code=True,
11
+ device="cpu",
12
+ config_kwargs={"use_memory_efficient_attention": False, "unpad_inputs": False}
13
+ )
14
+
15
+ docs_list = list(meme_attribute.values())
16
+
17
+ doc_embeddings = model.encode(docs_list)
18
+
19
+ embedded_dict = {key: embedding for key, embedding in zip(meme_attribute.keys(), doc_embeddings)}
20
+