amis-rpc-design/libs/amis/scripts/bot/embedding.py

14 lines
289 B
Python
Raw Normal View History

2023-10-07 19:42:30 +08:00
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)