| 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) |