File size: 730 Bytes
6666011 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 |
from sentence_transformers import SentenceTransformer
import litserve as ls
class EmbeddingAPI(ls.LitAPI):
def setup(self, device):
self.instruction = "Represent this sentence for searching relevant passages: "
self.model = SentenceTransformer('BAAI/bge-large-en-v1.5', device=device)
def decode_request(self, request):
return request["input"]
def predict(self, query):
return self.model.encode([self.instruction + query], normalize_embeddings=True)
def encode_response(self, output):
return {"embedding": output[0].tolist()}
if __name__ == "__main__":
api = EmbeddingAPI()
server = ls.LitServer(api, devices="cpu")
server.run(port=7860) |