Spaces:
Runtime error
Runtime error
Download app.py from hf1024/text-embeddings: direct link, hf CLI and curl.
- Browser
- Download file 13.1 kB
-
https://huggingface.co/spaces/hf1024/text-embeddings/resolve/main/app.py
- Command line
-
hf download hf://spaces/hf1024/text-embeddings/app.py
-
curl -L -o app.py https://huggingface.co/spaces/hf1024/text-embeddings/resolve/main/app.py
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", | |
| ) | |
| 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 [] | |
| 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) | |
| 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) | |
| 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") | |
| 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="/") | |