Spaces:
Sleeping
Sleeping
Upload 10 files
Browse files- Dockerfile +28 -0
- README (2).md +11 -0
- __init__.py +0 -0
- config.yml +25 -0
- get_embedding.py +130 -0
- gitattributes +35 -0
- gitignore +3 -0
- login.py +92 -0
- main.py +132 -0
- requirements.txt +57 -0
Dockerfile
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
FROM python:3.9-slim
|
| 2 |
+
|
| 3 |
+
# Install necessary system dependencies
|
| 4 |
+
RUN apt-get update && apt-get install -y \
|
| 5 |
+
build-essential \
|
| 6 |
+
libffi-dev \
|
| 7 |
+
libssl-dev \
|
| 8 |
+
&& rm -rf /var/lib/apt/lists/*
|
| 9 |
+
|
| 10 |
+
# Create and switch to a non-root user
|
| 11 |
+
RUN useradd -m -u 1000 user
|
| 12 |
+
USER user
|
| 13 |
+
|
| 14 |
+
# Set PATH so pip installs to user dir are available
|
| 15 |
+
ENV PATH="/home/user/.local/bin:$PATH"
|
| 16 |
+
|
| 17 |
+
WORKDIR /app
|
| 18 |
+
|
| 19 |
+
# Install Python dependencies
|
| 20 |
+
COPY --chown=user ./requirements.txt requirements.txt
|
| 21 |
+
RUN pip install --no-cache-dir --upgrade pip && \
|
| 22 |
+
pip install --no-cache-dir -r requirements.txt
|
| 23 |
+
|
| 24 |
+
# Copy all app files
|
| 25 |
+
COPY --chown=user . /app
|
| 26 |
+
|
| 27 |
+
# Run the app
|
| 28 |
+
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
|
README (2).md
ADDED
|
@@ -0,0 +1,11 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
title: Embedding
|
| 3 |
+
emoji: 🏢
|
| 4 |
+
colorFrom: gray
|
| 5 |
+
colorTo: red
|
| 6 |
+
sdk: docker
|
| 7 |
+
pinned: false
|
| 8 |
+
license: mit
|
| 9 |
+
---
|
| 10 |
+
|
| 11 |
+
Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
|
__init__.py
ADDED
|
File without changes
|
config.yml
ADDED
|
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
base_embedding_model: Qwen/Qwen3-Embedding-0.6B
|
| 2 |
+
ngrok_endpoint : ""
|
| 3 |
+
base_gpt_llm: gpt-4o-mini
|
| 4 |
+
base_ollama_llm: llama3.1:8b
|
| 5 |
+
hugging_face_url: Qwen/Qwen3-Embedding-0.6B
|
| 6 |
+
last_update_embedding_finetune: '2003-04-07T20:15:00.000Z'
|
| 7 |
+
last_update_faiss_push: '2025-03-02 02:25:23'
|
| 8 |
+
last_updates_preprocess_json: '2025-03-02 02:21:54'
|
| 9 |
+
server_url: 127.0.0.1
|
| 10 |
+
unique_titles:
|
| 11 |
+
- Conclusion
|
| 12 |
+
- Issue
|
| 13 |
+
- doc_citations
|
| 14 |
+
- doc_bench
|
| 15 |
+
- doc_title
|
| 16 |
+
- Analysis of the law
|
| 17 |
+
- pre_2
|
| 18 |
+
- Petitioner's Argument
|
| 19 |
+
- docsource_main
|
| 20 |
+
- Precedent Analysis
|
| 21 |
+
- pre_1
|
| 22 |
+
- Court's Reasoning
|
| 23 |
+
- Fact
|
| 24 |
+
- Respondent's Argument
|
| 25 |
+
- doc_author
|
get_embedding.py
ADDED
|
@@ -0,0 +1,130 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# get_embedding.py
|
| 2 |
+
from __future__ import annotations
|
| 3 |
+
|
| 4 |
+
import asyncio
|
| 5 |
+
import os
|
| 6 |
+
import yaml
|
| 7 |
+
import torch
|
| 8 |
+
from typing import List, Optional
|
| 9 |
+
|
| 10 |
+
from huggingface_hub import snapshot_download
|
| 11 |
+
from langchain_huggingface import HuggingFaceEmbeddings
|
| 12 |
+
|
| 13 |
+
# If you have an auth helper, we keep it
|
| 14 |
+
from login import HuggingFaceLogin
|
| 15 |
+
|
| 16 |
+
|
| 17 |
+
class EmbeddingFetcher:
|
| 18 |
+
"""
|
| 19 |
+
Async-friendly wrapper around a HuggingFace embedding model.
|
| 20 |
+
|
| 21 |
+
- Lazily initializes the model and downloads repo snapshot at app startup.
|
| 22 |
+
- Runs blocking HF / Torch operations in a worker thread via asyncio.to_thread.
|
| 23 |
+
- Thread-safe through an asyncio.Lock to prevent duplicate initialization.
|
| 24 |
+
"""
|
| 25 |
+
|
| 26 |
+
def __init__(self, config_path: str = "config.yml") -> None:
|
| 27 |
+
self._config_path = config_path
|
| 28 |
+
self._ready = False
|
| 29 |
+
self._init_lock = asyncio.Lock()
|
| 30 |
+
|
| 31 |
+
self._model: Optional[HuggingFaceEmbeddings] = None
|
| 32 |
+
self._local_model_path: Optional[str] = None
|
| 33 |
+
|
| 34 |
+
# Authenticate early (likely blocking), but off main thread in ensure_ready()
|
| 35 |
+
self._login = HuggingFaceLogin()
|
| 36 |
+
|
| 37 |
+
# Config defaults (overridden by config.yml)
|
| 38 |
+
self._repo_id: str = "SPAL0028/default-model"
|
| 39 |
+
self._normalize_embeddings: bool = False
|
| 40 |
+
|
| 41 |
+
# Device
|
| 42 |
+
self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| 43 |
+
|
| 44 |
+
# -----------------------
|
| 45 |
+
# Public properties
|
| 46 |
+
# -----------------------
|
| 47 |
+
@property
|
| 48 |
+
def model_id(self) -> str:
|
| 49 |
+
return self._local_model_path or self._repo_id
|
| 50 |
+
|
| 51 |
+
@property
|
| 52 |
+
def device_str(self) -> str:
|
| 53 |
+
return str(self._device)
|
| 54 |
+
|
| 55 |
+
# -----------------------
|
| 56 |
+
# Initialization
|
| 57 |
+
# -----------------------
|
| 58 |
+
async def ensure_ready(self) -> None:
|
| 59 |
+
"""
|
| 60 |
+
Idempotent async initializer. Safe to call multiple times.
|
| 61 |
+
"""
|
| 62 |
+
if self._ready:
|
| 63 |
+
return
|
| 64 |
+
|
| 65 |
+
async with self._init_lock:
|
| 66 |
+
if self._ready:
|
| 67 |
+
return
|
| 68 |
+
|
| 69 |
+
# 1) Login (may prompt environment-based auth); keep this off the main loop
|
| 70 |
+
await asyncio.to_thread(self._login.authenticate)
|
| 71 |
+
|
| 72 |
+
# 2) Load config
|
| 73 |
+
cfg = await asyncio.to_thread(self._load_config)
|
| 74 |
+
self._repo_id = cfg.get("hugging_face_url", self._repo_id)
|
| 75 |
+
self._normalize_embeddings = bool(cfg.get("normalize_embeddings", self._normalize_embeddings))
|
| 76 |
+
|
| 77 |
+
# 3) Download snapshot if needed (blocking IO) off main thread
|
| 78 |
+
self._local_model_path = await asyncio.to_thread(
|
| 79 |
+
snapshot_download,
|
| 80 |
+
self._repo_id
|
| 81 |
+
)
|
| 82 |
+
|
| 83 |
+
# 4) Build embeddings (blocking CPU-bound creation) off main thread
|
| 84 |
+
def _build_embeddings():
|
| 85 |
+
return HuggingFaceEmbeddings(
|
| 86 |
+
model_name=self._local_model_path,
|
| 87 |
+
model_kwargs={"device": self._device},
|
| 88 |
+
encode_kwargs={"normalize_embeddings": self._normalize_embeddings},
|
| 89 |
+
)
|
| 90 |
+
|
| 91 |
+
self._model = await asyncio.to_thread(_build_embeddings)
|
| 92 |
+
|
| 93 |
+
self._ready = True
|
| 94 |
+
|
| 95 |
+
# -----------------------
|
| 96 |
+
# Core API
|
| 97 |
+
# -----------------------
|
| 98 |
+
async def embed(self, texts: List[str] | str) -> List[List[float]]:
|
| 99 |
+
"""
|
| 100 |
+
Generate embeddings for a list of texts.
|
| 101 |
+
Ensures the model is initialized and runs blocking ops using a thread.
|
| 102 |
+
"""
|
| 103 |
+
await self.ensure_ready()
|
| 104 |
+
|
| 105 |
+
if isinstance(texts, str):
|
| 106 |
+
texts = [texts]
|
| 107 |
+
|
| 108 |
+
if not texts:
|
| 109 |
+
raise ValueError("No texts provided for embedding.")
|
| 110 |
+
|
| 111 |
+
# Defensive strip to avoid empty items sneaking through
|
| 112 |
+
sanitized = [t if isinstance(t, str) else str(t) for t in texts]
|
| 113 |
+
sanitized = [t.strip() for t in sanitized if t and t.strip()]
|
| 114 |
+
|
| 115 |
+
if not sanitized:
|
| 116 |
+
raise ValueError("All provided texts are empty after sanitization.")
|
| 117 |
+
|
| 118 |
+
# langchain_huggingface.HuggingFaceEmbeddings.embed_documents is blocking
|
| 119 |
+
vectors = await asyncio.to_thread(self._model.embed_documents, sanitized) # type: ignore[union-attr]
|
| 120 |
+
return vectors
|
| 121 |
+
|
| 122 |
+
# -----------------------
|
| 123 |
+
# Helpers
|
| 124 |
+
# -----------------------
|
| 125 |
+
def _load_config(self) -> dict:
|
| 126 |
+
if not os.path.exists(self._config_path):
|
| 127 |
+
# Keep behavior predictable if config is missing
|
| 128 |
+
return {}
|
| 129 |
+
with open(self._config_path, "r", encoding="utf-8") as f:
|
| 130 |
+
return yaml.safe_load(f) or {}
|
gitattributes
ADDED
|
@@ -0,0 +1,35 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
*.7z filter=lfs diff=lfs merge=lfs -text
|
| 2 |
+
*.arrow filter=lfs diff=lfs merge=lfs -text
|
| 3 |
+
*.bin filter=lfs diff=lfs merge=lfs -text
|
| 4 |
+
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
| 5 |
+
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
| 6 |
+
*.ftz filter=lfs diff=lfs merge=lfs -text
|
| 7 |
+
*.gz filter=lfs diff=lfs merge=lfs -text
|
| 8 |
+
*.h5 filter=lfs diff=lfs merge=lfs -text
|
| 9 |
+
*.joblib filter=lfs diff=lfs merge=lfs -text
|
| 10 |
+
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
| 11 |
+
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
| 12 |
+
*.model filter=lfs diff=lfs merge=lfs -text
|
| 13 |
+
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
| 14 |
+
*.npy filter=lfs diff=lfs merge=lfs -text
|
| 15 |
+
*.npz filter=lfs diff=lfs merge=lfs -text
|
| 16 |
+
*.onnx filter=lfs diff=lfs merge=lfs -text
|
| 17 |
+
*.ot filter=lfs diff=lfs merge=lfs -text
|
| 18 |
+
*.parquet filter=lfs diff=lfs merge=lfs -text
|
| 19 |
+
*.pb filter=lfs diff=lfs merge=lfs -text
|
| 20 |
+
*.pickle filter=lfs diff=lfs merge=lfs -text
|
| 21 |
+
*.pkl filter=lfs diff=lfs merge=lfs -text
|
| 22 |
+
*.pt filter=lfs diff=lfs merge=lfs -text
|
| 23 |
+
*.pth filter=lfs diff=lfs merge=lfs -text
|
| 24 |
+
*.rar filter=lfs diff=lfs merge=lfs -text
|
| 25 |
+
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
| 26 |
+
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
| 27 |
+
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
| 28 |
+
*.tar filter=lfs diff=lfs merge=lfs -text
|
| 29 |
+
*.tflite filter=lfs diff=lfs merge=lfs -text
|
| 30 |
+
*.tgz filter=lfs diff=lfs merge=lfs -text
|
| 31 |
+
*.wasm filter=lfs diff=lfs merge=lfs -text
|
| 32 |
+
*.xz filter=lfs diff=lfs merge=lfs -text
|
| 33 |
+
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
+
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
+
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
gitignore
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# .gitignore
|
| 2 |
+
.env
|
| 3 |
+
*.env
|
login.py
ADDED
|
@@ -0,0 +1,92 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# login.py
|
| 2 |
+
from __future__ import annotations
|
| 3 |
+
|
| 4 |
+
import os
|
| 5 |
+
import threading
|
| 6 |
+
import logging
|
| 7 |
+
from typing import Optional
|
| 8 |
+
|
| 9 |
+
from dotenv import load_dotenv
|
| 10 |
+
from huggingface_hub import login as hf_login
|
| 11 |
+
|
| 12 |
+
_logger = logging.getLogger(__name__)
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
class HuggingFaceLogin:
|
| 16 |
+
"""
|
| 17 |
+
Lightweight, idempotent wrapper around Hugging Face Hub login.
|
| 18 |
+
|
| 19 |
+
- Reads token from environment (.env supported).
|
| 20 |
+
- Safe to call authenticate() multiple times (thread-safe).
|
| 21 |
+
- Avoids writing to git credential store by default.
|
| 22 |
+
"""
|
| 23 |
+
|
| 24 |
+
_env_loaded = False
|
| 25 |
+
_env_lock = threading.Lock()
|
| 26 |
+
|
| 27 |
+
def __init__(
|
| 28 |
+
self,
|
| 29 |
+
token: Optional[str] = None,
|
| 30 |
+
*,
|
| 31 |
+
read_dotenv: bool = True,
|
| 32 |
+
add_to_git_credential: bool = False,
|
| 33 |
+
) -> None:
|
| 34 |
+
# Load env only once per process (thread-safe)
|
| 35 |
+
if read_dotenv:
|
| 36 |
+
with self._env_lock:
|
| 37 |
+
if not self._env_loaded:
|
| 38 |
+
load_dotenv()
|
| 39 |
+
self.__class__._env_loaded = True
|
| 40 |
+
|
| 41 |
+
# Accept explicit token or fall back to common env vars
|
| 42 |
+
self._token = token or self._read_token_from_env()
|
| 43 |
+
if not self._token:
|
| 44 |
+
raise ValueError(
|
| 45 |
+
"Hugging Face API token not found. "
|
| 46 |
+
"Set one of HF_TOKEN / HUGGINGFACEHUB_API_TOKEN / HUGGINGFACE_TOKEN in your environment "
|
| 47 |
+
"or pass token=... to HuggingFaceLogin()."
|
| 48 |
+
)
|
| 49 |
+
|
| 50 |
+
self._add_to_git_credential = add_to_git_credential
|
| 51 |
+
self._did_authenticate = False
|
| 52 |
+
self._auth_lock = threading.Lock()
|
| 53 |
+
|
| 54 |
+
# -----------------------
|
| 55 |
+
# Public API
|
| 56 |
+
# -----------------------
|
| 57 |
+
def authenticate(self) -> bool:
|
| 58 |
+
"""
|
| 59 |
+
Perform a one-time login to the Hugging Face Hub.
|
| 60 |
+
Safe to call multiple times (no-op after first success).
|
| 61 |
+
Returns True if authenticated in this call or previously.
|
| 62 |
+
"""
|
| 63 |
+
if self._did_authenticate:
|
| 64 |
+
return True
|
| 65 |
+
|
| 66 |
+
with self._auth_lock:
|
| 67 |
+
if self._did_authenticate:
|
| 68 |
+
return True
|
| 69 |
+
|
| 70 |
+
# Perform an in-process login; do not persist to git credentials by default
|
| 71 |
+
hf_login(
|
| 72 |
+
token=self._token,
|
| 73 |
+
add_to_git_credential=self._add_to_git_credential,
|
| 74 |
+
)
|
| 75 |
+
self._did_authenticate = True
|
| 76 |
+
_logger.info("✅ Authenticated with Hugging Face Hub.")
|
| 77 |
+
return True
|
| 78 |
+
|
| 79 |
+
@property
|
| 80 |
+
def is_authenticated(self) -> bool:
|
| 81 |
+
return self._did_authenticate
|
| 82 |
+
|
| 83 |
+
# -----------------------
|
| 84 |
+
# Helpers
|
| 85 |
+
# -----------------------
|
| 86 |
+
def _read_token_from_env(self) -> Optional[str]:
|
| 87 |
+
# Check common env var names in order of preference
|
| 88 |
+
for key in ("HF_TOKEN", "HUGGINGFACEHUB_API_TOKEN", "HUGGINGFACE_TOKEN"):
|
| 89 |
+
val = os.getenv(key)
|
| 90 |
+
if val and val.strip():
|
| 91 |
+
return val.strip()
|
| 92 |
+
return None
|
main.py
ADDED
|
@@ -0,0 +1,132 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# main.py
|
| 2 |
+
from __future__ import annotations
|
| 3 |
+
|
| 4 |
+
import asyncio
|
| 5 |
+
import time
|
| 6 |
+
from typing import List, Optional, Dict, Any
|
| 7 |
+
|
| 8 |
+
from fastapi import FastAPI, HTTPException, Depends
|
| 9 |
+
from fastapi.middleware.cors import CORSMiddleware
|
| 10 |
+
from pydantic import BaseModel, Field, validator
|
| 11 |
+
|
| 12 |
+
from get_embedding import EmbeddingFetcher
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
# -----------------------------
|
| 16 |
+
# Request / Response Schemas
|
| 17 |
+
# -----------------------------
|
| 18 |
+
class TextListRequest(BaseModel):
|
| 19 |
+
texts: List[str] = Field(..., description="List of strings to embed", min_items=1)
|
| 20 |
+
|
| 21 |
+
@validator("texts")
|
| 22 |
+
def non_empty_texts(cls, v: List[str]) -> List[str]:
|
| 23 |
+
if any((t is None or not isinstance(t, str) or t.strip() == "") for t in v):
|
| 24 |
+
raise ValueError("All items in 'texts' must be non-empty strings.")
|
| 25 |
+
return v
|
| 26 |
+
|
| 27 |
+
|
| 28 |
+
class EmbeddingResponse(BaseModel):
|
| 29 |
+
model_id: str
|
| 30 |
+
device: str
|
| 31 |
+
dims: int
|
| 32 |
+
count: int
|
| 33 |
+
elapsed_ms: float
|
| 34 |
+
embeddings: List[List[float]]
|
| 35 |
+
|
| 36 |
+
|
| 37 |
+
# -----------------------------
|
| 38 |
+
# App factory with lifespan
|
| 39 |
+
# -----------------------------
|
| 40 |
+
def create_app() -> FastAPI:
|
| 41 |
+
app = FastAPI(title="Embedding API", version="1.0.0")
|
| 42 |
+
|
| 43 |
+
# CORS: keep your original open policy (tighten in production)
|
| 44 |
+
app.add_middleware(
|
| 45 |
+
CORSMiddleware,
|
| 46 |
+
allow_origins=["*"],
|
| 47 |
+
allow_credentials=True,
|
| 48 |
+
allow_methods=["*"],
|
| 49 |
+
allow_headers=["*"],
|
| 50 |
+
)
|
| 51 |
+
|
| 52 |
+
# Global container for services
|
| 53 |
+
app.state.container: Dict[str, Any] = {}
|
| 54 |
+
app.state.init_lock = asyncio.Lock()
|
| 55 |
+
|
| 56 |
+
@app.on_event("startup")
|
| 57 |
+
async def on_startup() -> None:
|
| 58 |
+
# Initialize the EmbeddingFetcher once, asynchronously
|
| 59 |
+
async with app.state.init_lock:
|
| 60 |
+
if "embedder" not in app.state.container:
|
| 61 |
+
fetcher = EmbeddingFetcher()
|
| 62 |
+
# Build models / download snapshots off the main thread
|
| 63 |
+
await fetcher.ensure_ready()
|
| 64 |
+
app.state.container["embedder"] = fetcher
|
| 65 |
+
|
| 66 |
+
@app.on_event("shutdown")
|
| 67 |
+
async def on_shutdown() -> None:
|
| 68 |
+
# Nothing special is required, but the hook is here for future cleanup
|
| 69 |
+
pass
|
| 70 |
+
|
| 71 |
+
# ---------------
|
| 72 |
+
# Dependencies
|
| 73 |
+
# ---------------
|
| 74 |
+
def get_embedder() -> EmbeddingFetcher:
|
| 75 |
+
fetcher: Optional[EmbeddingFetcher] = app.state.container.get("embedder")
|
| 76 |
+
if fetcher is None:
|
| 77 |
+
# Defensive: if a request sneaks in before startup finishes
|
| 78 |
+
raise HTTPException(status_code=503, detail="Service not ready. Try again shortly.")
|
| 79 |
+
return fetcher
|
| 80 |
+
|
| 81 |
+
# ---------------
|
| 82 |
+
# Routes
|
| 83 |
+
# ---------------
|
| 84 |
+
@app.get("/", tags=["meta"])
|
| 85 |
+
async def home():
|
| 86 |
+
return {"status": "ok", "message": "Embedding service is running."}
|
| 87 |
+
|
| 88 |
+
@app.get("/healthz", tags=["meta"])
|
| 89 |
+
async def healthz():
|
| 90 |
+
# Lightweight health; could add a test encode if you want deeper checks
|
| 91 |
+
return {"status": "healthy"}
|
| 92 |
+
|
| 93 |
+
@app.post("/get-embedding/", response_model=EmbeddingResponse, tags=["embedding"])
|
| 94 |
+
async def get_embedding(request: TextListRequest, embedder: EmbeddingFetcher = Depends(get_embedder)):
|
| 95 |
+
# Offload embedding to the service (async wrapper over blocking HF/Torch calls)
|
| 96 |
+
start = time.perf_counter()
|
| 97 |
+
try:
|
| 98 |
+
vectors = await embedder.embed(request.texts)
|
| 99 |
+
except ValueError as ve:
|
| 100 |
+
raise HTTPException(status_code=400, detail=str(ve)) from ve
|
| 101 |
+
except Exception as e:
|
| 102 |
+
raise HTTPException(status_code=500, detail=f"Embedding failed: {e}") from e
|
| 103 |
+
|
| 104 |
+
elapsed_ms = (time.perf_counter() - start) * 1000.0
|
| 105 |
+
dims = len(vectors[0]) if vectors and len(vectors[0]) else 0
|
| 106 |
+
|
| 107 |
+
return EmbeddingResponse(
|
| 108 |
+
model_id=embedder.model_id,
|
| 109 |
+
device=embedder.device_str,
|
| 110 |
+
dims=dims,
|
| 111 |
+
count=len(vectors),
|
| 112 |
+
elapsed_ms=round(elapsed_ms, 3),
|
| 113 |
+
embeddings=vectors,
|
| 114 |
+
)
|
| 115 |
+
|
| 116 |
+
return app
|
| 117 |
+
|
| 118 |
+
|
| 119 |
+
app = create_app()
|
| 120 |
+
|
| 121 |
+
# Optional: run via `python main.py` in development
|
| 122 |
+
if __name__ == "__main__":
|
| 123 |
+
import uvicorn
|
| 124 |
+
|
| 125 |
+
uvicorn.run(
|
| 126 |
+
"main:app",
|
| 127 |
+
host="0.0.0.0",
|
| 128 |
+
port=8000,
|
| 129 |
+
reload=True, # Turn off in production
|
| 130 |
+
workers=1, # Use a process manager (e.g., gunicorn) to scale
|
| 131 |
+
log_level="info",
|
| 132 |
+
)
|
requirements.txt
ADDED
|
@@ -0,0 +1,57 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
annotated-types
|
| 2 |
+
anyio
|
| 3 |
+
certifi
|
| 4 |
+
charset-normalizer
|
| 5 |
+
click
|
| 6 |
+
colorama
|
| 7 |
+
fastapi
|
| 8 |
+
filelock
|
| 9 |
+
fsspec
|
| 10 |
+
h11
|
| 11 |
+
httpcore
|
| 12 |
+
httptools
|
| 13 |
+
httpx
|
| 14 |
+
huggingface-hub
|
| 15 |
+
idna
|
| 16 |
+
Jinja2
|
| 17 |
+
joblib
|
| 18 |
+
jsonpatch
|
| 19 |
+
jsonpointer
|
| 20 |
+
langchain-core
|
| 21 |
+
langchain-huggingface
|
| 22 |
+
langsmith
|
| 23 |
+
MarkupSafe
|
| 24 |
+
mpmath
|
| 25 |
+
networkx
|
| 26 |
+
numpy
|
| 27 |
+
orjson
|
| 28 |
+
packaging
|
| 29 |
+
pillow
|
| 30 |
+
pydantic
|
| 31 |
+
pydantic_core
|
| 32 |
+
python-dotenv
|
| 33 |
+
PyYAML
|
| 34 |
+
regex
|
| 35 |
+
requests
|
| 36 |
+
requests-toolbelt
|
| 37 |
+
safetensors
|
| 38 |
+
scikit-learn
|
| 39 |
+
scipy
|
| 40 |
+
sentence-transformers
|
| 41 |
+
sniffio
|
| 42 |
+
starlette
|
| 43 |
+
sympy
|
| 44 |
+
tenacity
|
| 45 |
+
threadpoolctl
|
| 46 |
+
tokenizers
|
| 47 |
+
torch
|
| 48 |
+
tqdm
|
| 49 |
+
transformers
|
| 50 |
+
typing-inspection
|
| 51 |
+
typing_extensions
|
| 52 |
+
urllib3
|
| 53 |
+
uvicorn
|
| 54 |
+
watchfiles
|
| 55 |
+
websockets
|
| 56 |
+
yml
|
| 57 |
+
zstandard
|