|
| 1 | +import json |
| 2 | +import os |
| 3 | +import numpy as np |
| 4 | +import onnxruntime as ort |
| 5 | +from http.server import HTTPServer, BaseHTTPRequestHandler |
| 6 | +from tokenizers import Tokenizer |
| 7 | + |
| 8 | +MODEL_DIR = os.environ.get("MODEL_DIR", "/opt/model") |
| 9 | +PORT = int(os.environ.get("EMBEDDING_PORT", "9222")) |
| 10 | + |
| 11 | +print(f"Loading tokenizer from {MODEL_DIR}...", flush=True) |
| 12 | +tokenizer = Tokenizer.from_file(os.path.join(MODEL_DIR, "tokenizer.json")) |
| 13 | +tokenizer.enable_padding(length=512) |
| 14 | +tokenizer.enable_truncation(max_length=512) |
| 15 | + |
| 16 | +print(f"Loading ONNX model from {MODEL_DIR}...", flush=True) |
| 17 | +session = ort.InferenceSession( |
| 18 | + os.path.join(MODEL_DIR, "model.onnx"), |
| 19 | + providers=["CPUExecutionProvider"] |
| 20 | +) |
| 21 | +print(f"Model loaded (ready to serve on port {PORT})", flush=True) |
| 22 | + |
| 23 | + |
| 24 | +class EmbedHandler(BaseHTTPRequestHandler): |
| 25 | + def do_POST(self): |
| 26 | + if self.path != "/embed": |
| 27 | + self._send(404, {"error": "not found"}) |
| 28 | + return |
| 29 | + try: |
| 30 | + length = int(self.headers.get("Content-Length", 0)) |
| 31 | + body = json.loads(self.rfile.read(length)) |
| 32 | + text = body.get("text", "") |
| 33 | + if not text: |
| 34 | + self._send(400, {"error": "text is required"}) |
| 35 | + return |
| 36 | + |
| 37 | + encoded = tokenizer.encode(text) |
| 38 | + input_ids = np.array([encoded.ids], dtype=np.int64) |
| 39 | + attention_mask = np.array([encoded.attention_mask], dtype=np.int64) |
| 40 | + token_type_ids = np.zeros_like(input_ids) |
| 41 | + |
| 42 | + outputs = session.run(None, { |
| 43 | + "input_ids": input_ids, |
| 44 | + "attention_mask": attention_mask, |
| 45 | + "token_type_ids": token_type_ids, |
| 46 | + }) |
| 47 | + |
| 48 | + cls_embedding = outputs[0][0][0].astype(float) |
| 49 | + norm = np.linalg.norm(cls_embedding) |
| 50 | + if norm > 0: |
| 51 | + cls_embedding = cls_embedding / norm |
| 52 | + |
| 53 | + self._send(200, {"embedding": cls_embedding.tolist()}) |
| 54 | + except Exception as e: |
| 55 | + self._send(500, {"error": str(e)}) |
| 56 | + |
| 57 | + def do_GET(self): |
| 58 | + if self.path == "/health": |
| 59 | + self._send(200, {"status": "ready"}) |
| 60 | + else: |
| 61 | + self._send(404, {"error": "not found"}) |
| 62 | + |
| 63 | + def _send(self, code, obj): |
| 64 | + self.send_response(code) |
| 65 | + self.send_header("Content-Type", "application/json") |
| 66 | + self.end_headers() |
| 67 | + self.wfile.write(json.dumps(obj).encode()) |
| 68 | + |
| 69 | + def log_message(self, fmt, *args): |
| 70 | + pass |
| 71 | + |
| 72 | + |
| 73 | +if __name__ == "__main__": |
| 74 | + server = HTTPServer(("0.0.0.0", PORT), EmbedHandler) |
| 75 | + print(f"Embedding server listening on port {PORT}", flush=True) |
| 76 | + server.serve_forever() |
0 commit comments