text-embeddings / app.py
hf1024's picture
Create app.py
0a72321 verified
Raw History Blame Contribute Delete
13.1 kB
# !/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@Author : Ailove
------------------------------------
@File : app.py.py
@CreateTime : 2026/7/24 13:54
------------------------------------
"""
import json
import os
import secrets
import tempfile
import time
from typing import Any, Tuple
import gradio as gr
import numpy as np
import spaces
import torch
from fastapi import Depends, FastAPI, Header, HTTPException, Request, status
from fastapi.exceptions import RequestValidationError
from fastapi.responses import JSONResponse
from pydantic import BaseModel, ConfigDict, Field
from transformers import AutoModel, AutoTokenizer
MODEL_NAME = "Qwen/Qwen3-Embedding-8B"
MODEL_ALIASES = {MODEL_NAME, "qwen3-embedding-8b"}
ALLOWED_DIMS = [4096, 2048, 1024, 768, 512, 256, 128]
MAX_DISPLAY_ROWS = 5
MAX_DISPLAY_VALUES = 16
MAX_BATCH_SIZE = int(os.getenv("MAX_BATCH_SIZE", "32"))
MAX_INPUT_CHARS = int(os.getenv("MAX_INPUT_CHARS", "20000"))
TOKENIZER = None
MODEL = None
def _load_model():
global TOKENIZER, MODEL
if MODEL is None:
TOKENIZER = AutoTokenizer.from_pretrained(MODEL_NAME)
MODEL = AutoModel.from_pretrained(MODEL_NAME, torch_dtype=torch.float16)
MODEL.eval()
if torch.cuda.is_available():
MODEL = MODEL.to("cuda")
def _l2_normalize(x: np.ndarray, eps: float = 1e-12) -> np.ndarray:
norms = np.linalg.norm(x, axis=1, keepdims=True)
return x / np.maximum(norms, eps)
def _mean_pool(last_hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
mask = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()
summed = torch.sum(last_hidden_state * mask, dim=1)
counts = torch.clamp(mask.sum(dim=1), min=1e-9)
return summed / counts
def _api_tokens() -> list[str]:
tokens = []
token_list = os.getenv("EMBEDDINGS_API_TOKENS", "")
single_token = os.getenv("EMBEDDINGS_API_TOKEN", "")
tokens.extend(token.strip() for token in token_list.split(",") if token.strip())
if single_token.strip():
tokens.append(single_token.strip())
return tokens
def _token_matches(provided_token: str, expected_tokens: list[str]) -> bool:
return bool(provided_token) and any(secrets.compare_digest(provided_token, token) for token in expected_tokens)
def _openai_error(
message: str,
status_code: int,
error_type: str = "invalid_request_error",
param: str | None = None,
code: str | None = None,
headers: dict[str, str] | None = None,
) -> JSONResponse:
return JSONResponse(
status_code=status_code,
headers=headers,
content={
"error": {
"message": message,
"type": error_type,
"param": param,
"code": code,
}
},
)
def _raise_openai_error(
message: str,
status_code: int,
error_type: str = "invalid_request_error",
param: str | None = None,
code: str | None = None,
headers: dict[str, str] | None = None,
) -> None:
raise HTTPException(
status_code=status_code,
detail={
"message": message,
"type": error_type,
"param": param,
"code": code,
},
headers=headers,
)
def verify_bearer_token(authorization: str | None = Header(default=None)) -> None:
expected_tokens = _api_tokens()
if not expected_tokens:
_raise_openai_error(
"Embedding API token is not configured.",
status.HTTP_401_UNAUTHORIZED,
code="invalid_api_key",
headers={"WWW-Authenticate": "Bearer"},
)
prefix = "Bearer "
if authorization is None or not authorization.startswith(prefix):
_raise_openai_error(
"Missing or invalid API key.",
status.HTTP_401_UNAUTHORIZED,
code="invalid_api_key",
headers={"WWW-Authenticate": "Bearer"},
)
provided_token = authorization[len(prefix) :].strip()
if not _token_matches(provided_token, expected_tokens):
_raise_openai_error(
"Invalid API key.",
status.HTTP_401_UNAUTHORIZED,
code="invalid_api_key",
headers={"WWW-Authenticate": "Bearer"},
)
def _parse_ui_texts(text_input: str) -> list[str]:
texts = [line.strip() for line in (text_input or "").splitlines() if line.strip()]
if not texts:
raise gr.Error("Please provide at least one non-empty line of text.")
return texts
def _validate_texts(texts: list[str]) -> None:
if not texts:
_raise_openai_error("Input must contain at least one non-empty string.", status.HTTP_400_BAD_REQUEST, param="input")
if len(texts) > MAX_BATCH_SIZE:
_raise_openai_error(
f"Batch size cannot exceed {MAX_BATCH_SIZE}.",
status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
param="input",
code="request_too_large",
)
for index, text in enumerate(texts):
if not text.strip():
_raise_openai_error(
"Input strings must not be empty.",
status.HTTP_400_BAD_REQUEST,
param=f"input[{index}]",
)
if len(text) > MAX_INPUT_CHARS:
_raise_openai_error(
f"Input strings cannot exceed {MAX_INPUT_CHARS} characters.",
status.HTTP_413_REQUEST_ENTITY_TOO_LARGE,
param=f"input[{index}]",
code="request_too_large",
)
@spaces.GPU(duration=300)
def _embed_batch(texts: list[str], dimension: int) -> tuple[list[list[float]], int, float]:
if dimension not in ALLOWED_DIMS:
raise ValueError(f"Dimension must be one of {ALLOWED_DIMS}.")
_load_model()
start = time.perf_counter()
with torch.no_grad():
batch = TOKENIZER(
texts,
padding=True,
truncation=True,
max_length=8192,
return_tensors="pt",
)
prompt_tokens = int(batch["attention_mask"].sum().item())
if torch.cuda.is_available():
batch = {k: v.to("cuda") for k, v in batch.items()}
outputs = MODEL(**batch)
pooled = _mean_pool(outputs.last_hidden_state, batch["attention_mask"])
full_embeddings = pooled.detach().cpu().numpy().astype(np.float32)
full_embeddings = _l2_normalize(full_embeddings)
truncated = full_embeddings[:, :dimension]
truncated = _l2_normalize(truncated)
elapsed = time.perf_counter() - start
return truncated.tolist(), prompt_tokens, elapsed
class EmbeddingsRequest(BaseModel):
model_config = ConfigDict(protected_namespaces=())
model: str = Field(default=MODEL_NAME)
input: str | list[str]
dimensions: int = Field(default=4096)
encoding_format: str = Field(default="float")
user: str | None = None
def _normalize_api_input(value: str | list[str]) -> list[str]:
if isinstance(value, str):
return [value]
if isinstance(value, list) and all(isinstance(item, str) for item in value):
return value
_raise_openai_error("Input must be a string or an array of strings.", status.HTTP_400_BAD_REQUEST, param="input")
return []
@spaces.GPU(duration=300)
def embed_texts(text_input: str, dimension: int) -> Tuple[str, str, str]:
if dimension not in ALLOWED_DIMS:
raise gr.Error(f"Dimension must be one of {ALLOWED_DIMS}.")
texts = _parse_ui_texts(text_input)
result_list, _prompt_tokens, elapsed = _embed_batch(texts, dimension)
display_preview = []
for row in result_list[:MAX_DISPLAY_ROWS]:
preview_row = row[:MAX_DISPLAY_VALUES]
if len(row) > MAX_DISPLAY_VALUES:
preview_row = preview_row + ["..."]
display_preview.append(preview_row)
preview_payload = {
"preview_embeddings": display_preview,
"shown_rows": min(len(result_list), MAX_DISPLAY_ROWS),
"total_rows": len(result_list),
"shown_values_per_row": MAX_DISPLAY_VALUES,
"actual_dimension": dimension,
"note": "Download the JSON file for full embeddings.",
}
stats = {
"model": MODEL_NAME,
"dimension": dimension,
"count": len(result_list),
"inference_time_seconds": round(elapsed, 4),
}
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False, encoding="utf-8") as f:
json.dump(
{
"model": MODEL_NAME,
"dimension": dimension,
"count": len(result_list),
"inference_time_seconds": elapsed,
"embeddings": result_list,
},
f,
ensure_ascii=False,
)
file_path = f.name
return json.dumps(preview_payload, indent=2), json.dumps(stats, indent=2), file_path
api_app = FastAPI(docs_url=None, redoc_url=None, openapi_url=None)
@api_app.exception_handler(HTTPException)
async def openai_http_exception_handler(_request: Request, exc: HTTPException) -> JSONResponse:
if isinstance(exc.detail, dict) and {"message", "type", "param", "code"}.issubset(exc.detail):
return JSONResponse(
status_code=exc.status_code,
headers=exc.headers,
content={"error": exc.detail},
)
return _openai_error(str(exc.detail), exc.status_code)
@api_app.exception_handler(RequestValidationError)
async def openai_validation_exception_handler(_request: Request, exc: RequestValidationError) -> JSONResponse:
del _request
first_error = exc.errors()[0] if exc.errors() else {}
location = first_error.get("loc", [])
param = ".".join(str(part) for part in location if part != "body") or None
message = first_error.get("msg", "Invalid request body.")
return _openai_error(str(message), status.HTTP_400_BAD_REQUEST, param=param, code="invalid_request")
@api_app.post("/v1/embeddings", dependencies=[Depends(verify_bearer_token)])
async def create_embeddings(payload: EmbeddingsRequest) -> dict[str, Any]:
if payload.model not in MODEL_ALIASES:
_raise_openai_error(
f"The model '{payload.model}' does not exist.",
status.HTTP_404_NOT_FOUND,
param="model",
code="model_not_found",
)
if payload.dimensions not in ALLOWED_DIMS:
_raise_openai_error(
f"Dimensions must be one of {ALLOWED_DIMS}.",
status.HTTP_400_BAD_REQUEST,
param="dimensions",
code="invalid_dimensions",
)
if payload.encoding_format != "float":
_raise_openai_error(
"Only encoding_format='float' is supported.",
status.HTTP_400_BAD_REQUEST,
param="encoding_format",
code="unsupported_encoding_format",
)
texts = _normalize_api_input(payload.input)
_validate_texts(texts)
embeddings, prompt_tokens, _elapsed = _embed_batch(texts, payload.dimensions)
return {
"object": "list",
"data": [
{
"object": "embedding",
"index": index,
"embedding": embedding,
}
for index, embedding in enumerate(embeddings)
],
"model": MODEL_NAME,
"usage": {
"prompt_tokens": prompt_tokens,
"total_tokens": prompt_tokens,
},
}
DESCRIPTION = """
High-quality text embeddings using **Qwen3-Embedding-8B** — top of MTEB text retrieval benchmarks.
Enter one text per line for batch embedding. L2-normalized, Matryoshka truncation to selected dimension.
Best for: semantic search, RAG, document retrieval, pure text similarity.
"""
FOOTER = (
"Built by [Xavier Fuentes](https://huggingface.co/xavier-fuentes) @ "
"[AI Enablement Academy](https://enablement.academy)"
)
with gr.Blocks(title="Text Embeddings - Qwen3-Embedding-8B") as demo:
gr.Markdown("# Text Embeddings - Qwen3-Embedding-8B")
gr.Markdown(DESCRIPTION)
with gr.Row():
text_input = gr.Textbox(
label="Input text (one per line for batch)",
lines=10,
placeholder="Type one sentence per line...",
)
dimension = gr.Dropdown(
choices=ALLOWED_DIMS,
value=4096,
label="Embedding dimension (Matryoshka)",
)
with gr.Row():
embed_btn = gr.Button("Generate Embeddings", variant="primary")
clear_btn = gr.Button("Clear")
embeddings_preview = gr.Code(label="Embeddings JSON Preview (truncated)", language="json")
stats_output = gr.Code(label="Stats", language="json")
download_file = gr.File(label="Download full embeddings JSON")
embed_btn.click(
fn=embed_texts,
inputs=[text_input, dimension],
outputs=[embeddings_preview, stats_output, download_file],
api_name=False,
)
clear_btn.click(fn=lambda: ("", "", None), outputs=[embeddings_preview, stats_output, download_file], api_name=False)
gr.Markdown(FOOTER)
app = gr.mount_gradio_app(api_app, demo.queue(), path="/")