| from pathlib import Path |
| import os |
| import json |
| import unicodedata |
| import numpy as np |
| import sentencepiece as spm |
| import torch |
| from fastapi import FastAPI, HTTPException |
| from pydantic import BaseModel |
|
|
| ROOT = Path(__file__).resolve().parent |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| MODEL = torch.jit.load(str(ROOT / "splinet_embedder.ts"), map_location=DEVICE).eval() |
| SP = spm.SentencePieceProcessor(model_file=str(ROOT / "spiece.model")) |
|
|
| PAD_ID = 3 |
| CLS_ID = 4 |
| SEP_ID = 5 |
| MAX_LENGTH = 512 |
| DEFAULT_BATCH = int(os.environ.get("SPLINET_BATCH_SIZE", "256" if DEVICE.type == "cuda" else "8")) |
|
|
| def normalize_text(text): |
| text = unicodedata.normalize("NFKC", text or "").lower() |
| return " ".join(text.split()) |
|
|
| def tokenize(texts): |
| ids = torch.full((len(texts), MAX_LENGTH), PAD_ID, dtype=torch.long) |
| mask = torch.zeros((len(texts), MAX_LENGTH), dtype=torch.long) |
| for i, text in enumerate(texts): |
| pieces = list(SP.encode(normalize_text(text), out_type=int))[:MAX_LENGTH - 2] |
| seq = [CLS_ID] + pieces + [SEP_ID] |
| n = len(seq) |
| ids[i, :n] = torch.tensor(seq, dtype=torch.long) |
| mask[i, :n] = 1 |
| return ids, mask |
|
|
| def encode(texts, batch_size=None): |
| if isinstance(texts, str): |
| texts = [texts] |
| if not texts: |
| return np.empty((0, 768), dtype=np.float32) |
| bs = max(1, int(batch_size or DEFAULT_BATCH)) |
| output = [] |
| start = 0 |
| while start < len(texts): |
| current = min(bs, len(texts) - start) |
| try: |
| ids, mask = tokenize(texts[start:start + current]) |
| ids = ids.to(DEVICE) |
| mask = mask.to(DEVICE) |
| with torch.inference_mode(): |
| if DEVICE.type == "cuda": |
| with torch.autocast(device_type="cuda", dtype=torch.float16): |
| z = MODEL(ids, mask) |
| else: |
| z = MODEL(ids, mask) |
| output.append(z.float().cpu().numpy()) |
| start += current |
| except torch.OutOfMemoryError: |
| if DEVICE.type != "cuda" or current == 1: |
| raise |
| torch.cuda.empty_cache() |
| bs = max(1, current // 2) |
| return np.concatenate(output, axis=0).astype(np.float32) |
|
|
| def similarity(text_a, text_b): |
| z = encode([text_a, text_b]) |
| return float(np.dot(z[0], z[1])) |
|
|
| class EmbedRequest(BaseModel): |
| texts: list[str] |
| batch_size: int | None = None |
|
|
| class SimilarityRequest(BaseModel): |
| text_a: str |
| text_b: str |
|
|
| app = FastAPI(title="SpliNet Embedder API", version="1.0.0") |
|
|
| @app.get("/health") |
| def health(): |
| return {"status": "ok", "device": str(DEVICE)} |
|
|
| @app.get("/model-info") |
| def model_info(): |
| with open(ROOT / "model_info.json", "r", encoding="utf-8") as f: |
| return json.load(f) |
|
|
| @app.post("/embed") |
| def embed_endpoint(req: EmbedRequest): |
| try: |
| z = encode(req.texts, req.batch_size) |
| return {"count": len(req.texts), "dimension": 768, "embeddings": z.tolist()} |
| except Exception as e: |
| raise HTTPException(status_code=500, detail=str(e)) |
|
|
| @app.post("/similarity") |
| def similarity_endpoint(req: SimilarityRequest): |
| try: |
| return {"cosine_similarity": similarity(req.text_a, req.text_b)} |
| except Exception as e: |
| raise HTTPException(status_code=500, detail=str(e)) |
|
|
| if __name__ == "__main__": |
| import uvicorn |
| uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "8000"))) |
|
|
