mirror of
https://gitee.com/baidu/amis.git
synced 2024-11-30 02:58:05 +08:00
14 lines
289 B
Python
14 lines
289 B
Python
|
import os
|
||
|
|
from sentence_transformers import SentenceTransformer, util
|
||
|
|
|
||
|
|
model_name = 'moka-ai/m3e-base'
|
||
|
|
|
||
|
|
if os.getenv('EMBEDDING_MODEL') != None:
|
||
|
|
model_name = os.getenv('EMBEDDING_MODEL')
|
||
|
|
|
||
|
|
model = SentenceTransformer(model_name)
|
||
|
|
|
||
|
|
|
||
|
|
def get_embedding(text):
|
||
|
|
return model.encode(text)
|