Christoph Haas
6101df3bfe
- Multi-stage Docker build (simplified to single-stage with pre-converted ONNX) - HTTP server with ONNX inference - API secret authentication - Uses pre-converted all-MiniLM-L6-v2 ONNX model from onnx-community - Image size: ~373 MB Generated by Mistral Vibe. Co-Authored-By: Mistral Vibe <vibe@mistral.ai>
100 lines
3.7 KiB
Python
100 lines
3.7 KiB
Python
#!/usr/bin/env python3
|
|
import json, os, sys, numpy as np
|
|
import onnxruntime as ort
|
|
from http.server import BaseHTTPRequestHandler, HTTPServer
|
|
|
|
# Load model at startup
|
|
session = ort.InferenceSession("/app/model.onnx")
|
|
|
|
# Load vocab from vocab.txt
|
|
vocab = {}
|
|
with open("/app/vocab.txt", "r", encoding="utf-8") as f:
|
|
for idx, token in enumerate(f):
|
|
token = token.strip()
|
|
vocab[token] = idx
|
|
|
|
# Token IDs from tokenizer_config.json
|
|
with open("/app/tokenizer_config.json", "r") as f:
|
|
tokenizer_config = json.load(f)
|
|
|
|
# Map token strings to IDs using vocab
|
|
cls_token_id = vocab.get(tokenizer_config["cls_token"], 0)
|
|
sep_token_id = vocab.get(tokenizer_config["sep_token"], 0)
|
|
pad_token_id = vocab.get(tokenizer_config["pad_token"], 0)
|
|
unk_token_id = vocab.get(tokenizer_config["unk_token"], 0)
|
|
|
|
def wordpiece_tokenize(text):
|
|
text = text.lower()
|
|
tokens = []
|
|
buffer = ""
|
|
for char in text:
|
|
if char.isspace():
|
|
if buffer:
|
|
tokens.append(buffer if buffer in vocab else "[UNK]")
|
|
buffer = ""
|
|
else:
|
|
buffer += char
|
|
if buffer:
|
|
tokens.append(buffer if buffer in vocab else "[UNK]")
|
|
return tokens
|
|
|
|
def encode(text, max_len=128):
|
|
token_ids = [vocab.get(t, unk_token_id) for t in wordpiece_tokenize(text)]
|
|
input_ids = [cls_token_id] + token_ids + [sep_token_id]
|
|
if len(input_ids) > max_len:
|
|
input_ids = input_ids[:max_len]
|
|
else:
|
|
input_ids = input_ids + [pad_token_id] * (max_len - len(input_ids))
|
|
attention_mask = [1] * len(input_ids)
|
|
token_type_ids = [0] * len(input_ids)
|
|
return input_ids, attention_mask, token_type_ids
|
|
|
|
print("Model loaded. Starting server on port 8080...")
|
|
|
|
API_SECRET = os.getenv("API_SECRET")
|
|
|
|
class VectorHandler(BaseHTTPRequestHandler):
|
|
def do_POST(self):
|
|
if self.path != "/vector":
|
|
self.send_error(404)
|
|
return
|
|
|
|
# Check API secret
|
|
if API_SECRET:
|
|
auth_header = self.headers.get("Authorization", "")
|
|
if auth_header != f"Bearer {API_SECRET}":
|
|
self.send_error(401, "Unauthorized - Invalid or missing API key")
|
|
return
|
|
|
|
content_length = int(self.headers.get("Content-Length", 0))
|
|
body = self.rfile.read(content_length)
|
|
try:
|
|
req = json.loads(body)
|
|
text = req.get("text", "")
|
|
if not text:
|
|
self.send_error(400, "Text is required")
|
|
return
|
|
input_ids, attention_mask, token_type_ids = encode(text)
|
|
input_ids_np = np.array([input_ids], dtype=np.int64)
|
|
attention_mask_np = np.array([attention_mask], dtype=np.int64)
|
|
token_type_ids_np = np.array([token_type_ids], dtype=np.int64)
|
|
outputs = session.run(
|
|
["last_hidden_state"],
|
|
{"input_ids": input_ids_np, "attention_mask": attention_mask_np, "token_type_ids": token_type_ids_np}
|
|
)
|
|
token_embeddings = outputs[0][0]
|
|
input_mask_expanded = np.expand_dims(attention_mask_np[0], axis=-1)
|
|
sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=0)
|
|
sum_mask = np.maximum(np.sum(input_mask_expanded, axis=0), 1e-9)
|
|
sentence_embedding = (sum_embeddings / sum_mask).tolist()
|
|
resp = json.dumps({"vector": sentence_embedding})
|
|
self.send_response(200)
|
|
self.send_header("Content-Type", "application/json")
|
|
self.end_headers()
|
|
self.wfile.write(resp.encode())
|
|
except Exception as e:
|
|
self.send_error(500, str(e))
|
|
def log_message(self, format, *args):
|
|
pass # Suppress logs
|
|
|
|
HTTPServer(("0.0.0.0", int(os.getenv("PORT", 8080))), VectorHandler).serve_forever()
|