Spaces:
Sleeping
Sleeping
| import os | |
| from typing import Literal | |
| from fastapi import Depends, FastAPI, Header, HTTPException, status | |
| from pydantic import BaseModel, Field | |
| from sentence_transformers import SentenceTransformer | |
| MODEL_NAME = os.getenv("EMBEDDING_MODEL_NAME", "BAAI/bge-base-en-v1.5") | |
| SERVICE_TOKEN = os.getenv("EMBEDDING_SERVICE_TOKEN", "") | |
| PASSAGE_PREFIX = "Represent this passage for searching relevant passages: " | |
| QUERY_PREFIX = "Represent this sentence for searching relevant passages: " | |
| app = FastAPI(title="CognifyAI Embedding Service") | |
| model = SentenceTransformer(MODEL_NAME) | |
| class EmbedRequest(BaseModel): | |
| text: str = Field(min_length=1) | |
| kind: Literal["passage", "query"] = "passage" | |
| class EmbedBatchRequest(BaseModel): | |
| texts: list[str] = Field(min_length=1) | |
| kind: Literal["passage", "query"] = "passage" | |
| def require_token(authorization: str = Header(default="")) -> None: | |
| if not SERVICE_TOKEN: | |
| return | |
| expected = f"Bearer {SERVICE_TOKEN}" | |
| if authorization != expected: | |
| raise HTTPException( | |
| status_code=status.HTTP_401_UNAUTHORIZED, | |
| detail="Invalid embedding service token", | |
| ) | |
| def prefix_text(text: str, kind: str) -> str: | |
| if kind == "query": | |
| return f"{QUERY_PREFIX}{text}" | |
| return f"{PASSAGE_PREFIX}{text}" | |
| def health_check() -> dict: | |
| return {"status": "ok", "model": MODEL_NAME} | |
| def embed(payload: EmbedRequest) -> dict: | |
| embedding = model.encode( | |
| prefix_text(payload.text, payload.kind), | |
| normalize_embeddings=True, | |
| ) | |
| return {"embedding": embedding.tolist(), "model": MODEL_NAME} | |
| def embed_batch(payload: EmbedBatchRequest) -> dict: | |
| prefixed_texts = [prefix_text(text, payload.kind) for text in payload.texts] | |
| embeddings = model.encode( | |
| prefixed_texts, | |
| normalize_embeddings=True, | |
| show_progress_bar=False, | |
| ) | |
| return {"embeddings": embeddings.tolist(), "model": MODEL_NAME} | |