amis-rpc-design/libs/amis/scripts/bot/embedding.py
2023-10-07 19:42:30 +08:00

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)