Upload 45 files

.env.example

@@ -0,0 +1,40 @@
+# .env.example - copy to .env and edit
+
+# -----------------------------
+# Security / Access
+# -----------------------------
+# Static API key used to guard endpoints (REQUIRED )
+EMBED_API_KEY=MY_SECURE_KEY
+# Private hugging face token (REQUIRED)
+HUGGING_FACE_TOKEN=your_free_token_with_repo_read_access
+
+
+# -----------------------------
+# Model & Device Settings
+# -----------------------------
+# HuggingFace model identifier
+MODEL_NAME=Faysal4200/bge-m3-private
+
+# Target device: "cuda" for GPU, "cpu" for CPU, or leave empty for auto-detection
+DEVICE=
+
+# How many texts to process in one model forward pass
+BATCH_SIZE=8
+
+# Maximum input sequence length (BGE-M3 supports up to 8192)
+MAX_LENGTH=1024
+
+# -----------------------------
+# Server & Reliability
+# -----------------------------
+RETRY_ATTEMPTS=3
+RETRY_BACKOFF_SECONDS=2.0
+HOST=0.0.0.0
+PORT=7860
+WORKERS=1
+
+# -----------------------------
+# Logging (Premium Style)
+# -----------------------------
+# Log level: DEBUG, INFO, WARNING, ERROR, CRITICAL
+LOG_LEVEL=INFO

.gitignore

@@ -0,0 +1,2 @@
+# env with credential removed 
+.env

Dockerfile

@@ -0,0 +1,37 @@
+# Dockerfile - GPU & CPU friendly image
+# Uses NVIDIA CUDA runtime as base for GPU support. 
+# Falls back to CPU automatically if no GPU is detected by the application.
+
+FROM nvidia/cuda:12.1.1-runtime-ubuntu22.04
+
+# Prevent interactive prompts during apt-get
+ENV DEBIAN_FRONTEND=noninteractive
+ENV PYTHONDONTWRITEBYTECODE=1
+ENV PYTHONUNBUFFERED=1
+
+WORKDIR /app
+
+# Install Python and minimal system dependencies
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    python3.11 \
+    python3-pip \
+    python3.11-dev \
+    build-essential \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set python3.11 as the default python
+RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1 \
+    && update-alternatives --set python3 /usr/bin/python3.11 \
+    && python3 -m pip install --upgrade pip
+
+# Copy requirements and install
+# We install torch with cuda 12.1 index for explicit GPU support, but it works on CPU too.
+COPY ./requirements.txt /app/requirements.txt
+RUN pip install --no-cache-dir -r /app/requirements.txt --extra-index-url https://download.pytorch.org/whl/cu121
+
+# Copy app code
+COPY ./app /app/app
+
+EXPOSE 7860
+
+CMD ["python3", "-m", "app.main"]

README.md

@@ -1,10 +1,252 @@
----
-title: Embedding Server
-emoji: 🚀
-colorFrom: yellow
-colorTo: blue
-sdk: docker
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# BGE-M3 Multilingual Embedding Server
+
+A high-performance, robust, and production-ready embedding server utilizing the **BGE-M3** model. This server is designed to provide dense, multilingual embeddings with features like batch processing, automatic GPU/CPU detection, and strict security guards.
+
+---
+
+## 🚀 Key Features
+
+- **State-of-the-Art Model**: Uses `BAAI/bge-m3` for high-quality multilingual embeddings. U can Change it also. with any huggigface model name just make sure it gives 1024 size vector embedding for consistency.
+- **FastAPI Powered**: Asynchronous, high-performance API endpoints.
+- **Intelligent Device Management**: Automatically detects NVIDIA GPUs (CUDA) or falls back to CPU.
+- **Robust Batching**: Efficiently processes large lists of text with automatic chunking and OOM (Out of Memory) recovery.
+- **Strict Security**: Built-in guard for API Key enforcement.
+- **Premium Logging**: Beautiful, colorized, and structured logs for easy debugging and monitoring.
+- **Reliability**: Integrated retry logic for model loading and processing.
+
+---
+
+## 🛠️ Tech Stack
+
+- **Language**: Python 3.11+
+- **Web Framework**: FastAPI
+- **Deep Learning**: PyTorch & HuggingFace Transformers
+- **Validation**: Pydantic V2 & Pydantic Settings
+- **Logging**: Colorlog
+- **Process Manager**: Uvicorn
+
+---
+
+## 📥 Installation
+
+It is highly recommended to use a virtual environment to keep dependencies local to the project.
+
+1.  **Clone the project** and navigate to the directory.
+2.  **Create a virtual environment**:
+    ```cmd
+    python -m venv venv
+    ```
+3.  **Activate the environment**:
+    - Windows: `venv\Scripts\activate`
+    - Linux/Mac: `source venv/bin/activate`
+4.  **Install dependencies**:
+    ```cmd
+    pip install -r requirements.txt
+    ```
+
+---
+
+## ⚙️ Configuration
+
+Copy the `.env.example` file to `.env` and configure your settings. The application uses these variables to control everything from security to model performance.
+
+```env
+# -----------------------------
+# Security / Access
+# -----------------------------
+EMBED_API_KEY=your_secret_key_here
+HUGGING_FACE_TOKEN=your_free_token_with_repo_read_access
+
+# -----------------------------
+# Model & Device Settings
+# -----------------------------
+MODEL_NAME=BAAI/bge-m3       # Best if u copy this repo to your personal repo and use updated model name
+DEVICE=                      # Leave empty for auto-detection, or use "cuda" / "cpu"
+BATCH_SIZE=8                 # Adjust based on your VRAM/RAM
+MAX_LENGTH=1024              # Max token length
+
+# -----------------------------
+# Server & Reliability
+# -----------------------------
+RETRY_ATTEMPTS=3             # Retries for model loading/processing
+RETRY_BACKOFF_SECONDS=2.0    # Wait time between retries
+HOST=0.0.0.0
+PORT=8000
+WORKERS=1                    # Number of Uvicorn workers
+LOG_LEVEL=INFO               # DEBUG, INFO, WARNING, ERROR, CRITICAL
+```
+
+---
+
+## 🏁 Running Locally
+
+Start the server using Python's module runner:
+
+```cmd
+python -m app.main
+```
+
+The server will automatically log the detected device (GPU/CPU) and start listening on the configured port.
+
+---
+
+## 🐳 Docker Deployment
+
+You can deploy the server easily using Docker and Docker Compose. This is the recommended way for production.
+
+### 1. Build and Run
+
+```bash
+docker-compose up --build -d
+```
+
+This will start the server on the port defined in your `.env` (default `8000`).
+# Note: After building the image if u want to run the image without code just use this command
+# docker run -d --restart unless-stopped --name embed-server -p 8000:8000 --env-file .env vector-embedding-server-embed-server:latest
+# here vector-embedding-server-embed-server:latest is the image name which is build
+# .env should have the Configuration specified in the .env.example file
+
+### 2. GPU Acceleration (NVIDIA)
+
+To use your NVIDIA GPU with Docker, you must have the **NVIDIA Container Toolkit** installed on your host system.
+
+1.  Open `docker-compose.yml`.
+2.  Uncomment the `deploy` section at the bottom of the file:
+    ```yaml
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
+    ```
+3.  Restart the container:
+    ```bash
+    docker-compose up -d --force-recreate
+    ```
+
+The server logs will confirm if `cuda` is being used.
+
+---
+
+## 🚀 Run Docker Image (Without Source Code)
+
+After building the image, run the container using:
+```
+docker run -d --restart unless-stopped --name embed-server -p 8000:8000 --env-file .env vector-embedding-server-embed-server:latest
+```
+# Requirements
+
+- The Docker image vector-embedding-server-embed-server:latest must already be built.
+- A .env file must exist in the same directory.
+
+# Environment Setup
+
+Create your .env file based on:
+```
+.env.example
+```
+Example:
+```
+cp .env.example .env
+```
+Update the values inside .env before running the container.
+
+# Check Running Container
+```
+docker ps
+```
+You should see embed-server listed as running.
+
+---
+
+## 📡 API Endpoints
+
+### 1. Health Check
+
+`GET /embed/health`  
+Returns the server status and whether the model is loaded.
+
+### 2. Single Embedding
+
+`POST /embed/single`  
+Generates an embedding for a single text string.
+
+### 3. Batch Embedding
+
+`POST /embed/batch`  
+Generates embeddings for a list of strings efficiently.
+
+---
+
+## 🧪 Documentation & Testing
+
+Once the server is running, you can access the interactive API documentation at:
+
+- **Swagger UI**: `http://localhost:7860/docs`
+- **Redoc**: `http://localhost:7860/redoc`
+
+### Test with CURL
+
+**Note**: The server requires `x-api-key` header as configured in your `.env`.
+
+#### Single Embedding Request:
+
+```bash
+curl -X POST http://localhost:7860/embed/single \
+     -H "Content-Type: application/json" \
+     -H "x-api-key: MY_SECURE_KEY" \
+     -d '{"text": "Hello world, this is a test."}'
+```
+
+#### Batch Embedding Request:
+
+```bash
+curl -X POST http://localhost:7860/embed/batch \
+     -H "Content-Type: application/json" \
+     -H "x-api-key: MY_SECURE_KEY" \
+     -d '{
+       "texts": [
+         "First sentence to embed.",
+         "BGE-M3 handles multilingual text very well."
+       ]
+     }'
+```
+
+---
+
+## 🔧 Troubleshooting
+
+### `OSError: Could not create safetensors conversion PR`
+
+If you see this error in your logs, it means the `transformers` library is trying to auto-convert your PyTorch model to SafeTensors format but failing (often due to repo permissions or structure).
+
+**Solution:**
+
+We have included a script to manually convert and upload the SafeTensors model to your Hugging Face repository.
+
+**Prerequisites:**
+1.  **Write Access**: You must have **write access** to the Hugging Face repository defined in `MODEL_NAME`.
+2.  **Environment Variable**: Ensure your `HUGGING_FACE_TOKEN` in `.env` has **write permissions**.
+
+**Instructions:**
+Run this script **ONLY ONCE** and **ONLY** if you are facing the error:
+
+```bash
+python -m scripts.convert_safetensors
+```
+
+This will:
+1.  Load your current model.
+2.  Convert it to SafeTensors format locally.
+3.  Verify the converted model loads correctly.
+4.  Upload the `model.safetensors` file to your repository.
+
+Once done, the "safetensors not found" error will disappear, and your model loading speed will improve.
+
+---
+
+## 👨‍💻 Developed By
+
+**FAYSAL AHMMED**

app/__init__.py

@@ -0,0 +1,2 @@
+# BGE-M3 Multilingual Embedding Server
+__version__ = "1.0.0"

app/auth.py

@@ -0,0 +1,23 @@
+# app/auth.py
+"""
+Authorization dependency.
+"""
+
+from fastapi import Request, HTTPException, status, Header
+from .config import settings
+from .logger import logger
+
+
+async def api_key_guard(
+    request: Request,
+    x_api_key: str = Header(..., alias="x-api-key", description="API Key for authentication"),
+) -> None:
+    # Check API key
+    # FastAPI's Header(...) ensures the key is present. behavior if missing is 422.
+    # We still check the value.
+    if x_api_key != settings.EMBED_API_KEY:
+        logger.warning(f"Unauthorized: Invalid API key. Provided: {x_api_key[:4] if x_api_key else 'None'}...")
+        raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Invalid API key")
+
+    logger.debug(f"Auth passed for request from {request.client.host if request.client else 'unknown'}")
+

app/config.py

@@ -0,0 +1,72 @@
+# app/config.py
+"""
+Application configuration using Pydantic BaseSettings.
+All configuration MUST be provided via environment variables or the .env file.
+"""
+
+import json
+from pydantic import Field, field_validator
+from pydantic_settings import BaseSettings, SettingsConfigDict
+from typing import List, Optional, Any
+
+
+class Settings(BaseSettings):
+    # -----------------------------
+    # Security / Access
+    # -----------------------------
+    EMBED_API_KEY: str = Field(..., validation_alias="EMBED_API_KEY")
+    HUGGING_FACE_TOKEN: str = Field(..., validation_alias="HUGGING_FACE_TOKEN")
+    
+
+
+    # -----------------------------
+    # Model & Device Settings
+    # -----------------------------
+    MODEL_NAME: str = Field(..., validation_alias="MODEL_NAME")
+    DEVICE: Optional[str] = Field(..., validation_alias="DEVICE") # Can be empty in .env for auto-detection
+    BATCH_SIZE: int = Field(8, validation_alias="BATCH_SIZE")
+    MAX_LENGTH: int = Field(1024, validation_alias="MAX_LENGTH")
+
+    # -----------------------------
+    # Server & Reliability
+    # -----------------------------
+    RETRY_ATTEMPTS: int = Field(3, validation_alias="RETRY_ATTEMPTS")
+    RETRY_BACKOFF_SECONDS: float = Field(2.0, validation_alias="RETRY_BACKOFF_SECONDS")
+    HOST: str = Field("0.0.0.0", validation_alias="HOST")
+    PORT: int = Field(7860, validation_alias="PORT")
+    WORKERS: int = Field(1, validation_alias="WORKERS")
+
+    # -----------------------------
+    # Logging
+    # -----------------------------
+    LOG_LEVEL: str = Field("INFO", validation_alias="LOG_LEVEL")
+
+    # Pydantic Settings Configuration
+    model_config = SettingsConfigDict(
+        env_file=".env",
+        env_file_encoding="utf-8",
+        extra="ignore"
+    )
+
+
+
+
+# Load settings (Strict mode: will raise error if any key is missing)
+try:
+    settings = Settings()
+    # ------------------------------------------------------------------
+    # Explicitly set HF environment variables so all HF-based libraries 
+    # (transformers, huggingface_hub, etc.) pick up the token globally.
+    # If this not done then hugging face token will not work
+    # ------------------------------------------------------------------
+    import os
+    if settings.HUGGING_FACE_TOKEN:
+        os.environ["HF_TOKEN"] = settings.HUGGING_FACE_TOKEN
+        os.environ["HUGGINGFACE_HUB_TOKEN"] = settings.HUGGING_FACE_TOKEN
+        
+except Exception as e:
+    import sys
+    print(f"\n[CRITICAL ERROR] Configuration failed to load from environment/.env:")
+    print(f"Missing or invalid keys: {e}")
+    print("\nPlease ensure your .env file is complete according to .env.example\n")
+    sys.exit(1)

app/controllers/__init__.py

@@ -0,0 +1 @@
+from .embed_controller import router as embed_router

app/controllers/embed_controller.py

@@ -0,0 +1,54 @@
+"""
+API router for embedding endpoints.
+"""
+
+from fastapi import APIRouter, Depends, HTTPException, status
+from typing import Any, Dict
+
+from ..models.request_models import SingleRequest, BatchRequest
+from ..services.async_embeddings_service import AsyncEmbeddingsService
+from ..auth import api_key_guard
+from ..logger import logger
+
+router = APIRouter(prefix="/embed", tags=["embed"])
+
+# service is a module-level singleton
+service = AsyncEmbeddingsService()
+
+
+@router.post("/single")
+async def embed_single(req: SingleRequest, _: Any = Depends(api_key_guard)):
+    """Generate embedding for one text."""
+    try:
+        emb = await service.generate_single_embedding(req.text)
+        return {"embedding": emb}
+    except ValueError as ve:
+        # Warning since it's a client usage error (400)
+        logger.warning(f"Validation Error: {ve}")
+        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(ve))
+    except Exception as e:
+        # Error level since it's a server failure (500)
+        logger.error(f"Failed /embed/single request")
+        # Traceback will be handled by global exception handler
+        raise 
+
+
+@router.post("/batch")
+async def embed_batch(req: BatchRequest, _: Any = Depends(api_key_guard)):
+    """Efficient batch embedding endpoint."""
+    try:
+        embeddings = await service.generate_batch_embeddings(req.texts)
+        return {"embeddings": embeddings}
+    except Exception as e:
+        logger.error(f"Failed /embed/batch request for {len(req.texts)} items")
+        raise 
+
+
+@router.get("/health")
+async def health_check(_: Any = Depends(api_key_guard)) -> Dict[str, Any]:
+    """Verify server status and model readiness."""
+    return {
+        "status": "healthy", 
+        "model_loaded": service.is_ready,
+        "device": str(service.device) if service.is_ready else "not_loaded"
+    }

app/logger.py

@@ -0,0 +1,59 @@
+# app/logger.py
+"""
+Structured, colorized logger used across the app.
+
+- Uses colorlog to add color to log levels in console.
+- Central logger configuration prevents duplicate handlers.
+"""
+
+import logging
+import sys
+from typing import Optional
+
+import colorlog
+from .config import settings
+
+
+def get_logger(name: Optional[str] = None) -> logging.Logger:
+    logger = logging.getLogger(name)
+    if logger.handlers:
+        # already configured
+        return logger
+
+    # Get level from settings or default to INFO
+    log_level_str = settings.LOG_LEVEL.upper()
+    level = getattr(logging, log_level_str, logging.INFO)
+
+    # Console Handler with Color
+    handler = colorlog.StreamHandler(stream=sys.stdout)
+    handler.setLevel(level)
+
+    # Enhanced Formatter with better spacing and colors
+    formatter = colorlog.ColoredFormatter(
+        "%(log_color)s%(levelname)-8s%(reset)s | %(log_color)s%(asctime)s%(reset)s | %(cyan)s%(name)s%(reset)s | %(message)s",
+        datefmt="%H:%M:%S",
+        log_colors={
+            "DEBUG": "blue",
+            "INFO": "green",
+            "WARNING": "yellow",
+            "ERROR": "bold_red",
+            "CRITICAL": "bold_red,bg_white",
+        },
+        secondary_log_colors={
+            "message": {
+                "ERROR": "red",
+                "CRITICAL": "red",
+            }
+        },
+        style="%"
+    )
+    handler.setFormatter(formatter)
+
+    logger.setLevel(level)
+    logger.addHandler(handler)
+    logger.propagate = False
+    return logger
+
+
+# Global app logger
+logger = get_logger("bge-server")

app/main.py

@@ -0,0 +1,106 @@
+# app/main.py
+"""
+Application entrypoint.
+
+- Configures FastAPI app and CORS middleware.
+- Includes embed router.
+- Starts uvicorn when run as __main__.
+"""
+
+import torch
+from contextlib import asynccontextmanager
+import uvicorn
+from fastapi import FastAPI, Request
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+
+from .config import settings
+from .logger import logger
+from .controllers.embed_controller import router as embed_router, service as embedding_service
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Startup logic
+    gpu_available = torch.cuda.is_available()
+    device_to_use = settings.DEVICE or ("cuda" if gpu_available else "cpu")
+    
+    logger.info("┌────────────────────────────────────────┐")
+    logger.info("│          Embedding Server Start        │")
+    logger.info("└────────────────────────────────────────┘")
+    
+    if gpu_available:
+        logger.info(f"GPU Detected: [bold]{torch.cuda.get_device_name(0)}[/bold]")
+    else:
+        logger.warning("GPU not detected, falling back to CPU.")
+        
+    logger.info(f"Active Device: {device_to_use}")
+    logger.info(f"Serving on: {settings.HOST}:{settings.PORT}")
+
+    # Eagerly load the model to RAM/VRAM on startup
+    try:
+        await embedding_service.load_model()
+    except Exception as e:
+        logger.error("CRITICAL: Failed to load model during startup.")
+        logger.exception(e)
+        # We don't exit here to allow the health check to remain accessible, 
+        # but the model status will be 'unloaded'.
+    
+    yield
+    
+    # Shutdown logic
+    logger.info("Shutting down server...")
+
+app = FastAPI(
+    title="BGE-M3 Multilingual Embedding Server",
+    summary="High-performance API for BAAI/bge-m3 embeddings.",
+    description="""
+        ## Features
+        - **Multilingual Support**: Embed text in 100+ languages.
+        - **GPU Acceleration**: Automatically uses CUDA if available.
+        - **Batch Processing**: Efficiently handle lists of text.
+        - **Security**: Strict API Key protection.
+
+        ## Usage
+        - **Required Header**: `x-api-key` must be provided for all `POST` / `GET` requests.
+        - **Interactive Docs**: Use the **Try it out** button to test endpoints directly from the browser.
+        """,
+    version="1.0.0",
+    lifespan=lifespan,
+    docs_url="/docs",  
+    redoc_url="/redoc" 
+)
+
+# Global Exception Handler for unhandled errors
+@app.exception_handler(Exception)
+async def global_exception_handler(request: Request, exc: Exception):
+    logger.error(f"Unhandled Exception: {type(exc).__name__}")
+    logger.error(f"Path: {request.url.path}")
+    logger.exception(exc) # This will log the full traceback in red
+    return JSONResponse(
+        status_code=500,
+        content={"detail": "An internal server error occurred.", "type": type(exc).__name__},
+    )
+
+# CORS middleware - allow all (security handled by API Key guard)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["GET", "POST"],
+    allow_headers=["*"],
+)
+
+import time
+@app.middleware("http")
+async def add_process_time_header(request: Request, call_next):
+    start_time = time.perf_counter()
+    response = await call_next(request)
+    process_time = (time.perf_counter() - start_time) * 1000
+    logger.info(f"Request: {request.method} {request.url.path} - Completed in {process_time:.2f}ms")
+    return response
+
+app.include_router(embed_router)
+
+
+if __name__ == "__main__":
+    uvicorn.run("app.main:app", host=settings.HOST, port=settings.PORT, workers=settings.WORKERS)

app/models/__init__.py

@@ -0,0 +1 @@
+from .request_models import SingleRequest, BatchRequest

app/models/request_models.py

@@ -0,0 +1,25 @@
+# app/models/request_models.py
+"""
+Pydantic request models for the API endpoints.
+"""
+
+from typing import List
+from pydantic import BaseModel, Field
+
+
+class SingleRequest(BaseModel):
+    text: str = Field(
+        ..., 
+        description="The text content to be embedded.",
+        min_length=1,
+        examples=["Hello, this is a test sentence."]
+    )
+
+
+class BatchRequest(BaseModel):
+    texts: List[str] = Field(
+        ..., 
+        description="List of text strings to embed.",
+        min_items=1,
+        examples=[["Hello world", "This is another sentence"]]
+    )

app/services/__init__.py

@@ -0,0 +1 @@
+from .embeddings_service import EmbeddingsService

app/services/async_embeddings_service.py

@@ -0,0 +1,50 @@
+"""
+AsyncEmbeddingsService:
+- Wraps the synchronous EmbeddingsService to provide async interface
+- Uses asyncio.to_thread to offload blocking model operations to a thread pool
+- Manages concurrency with asyncio.Lock
+"""
+
+import asyncio
+from typing import List, Optional
+from .embeddings_service import EmbeddingsService
+
+class AsyncEmbeddingsService:
+    def __init__(self, model_name: Optional[str] = None):
+        # We compose the synchronous service
+        self._sync_service = EmbeddingsService(model_name)
+        self._lock = asyncio.Lock()
+
+    @property
+    def is_ready(self) -> bool:
+        return self._sync_service.is_ready
+
+    @property
+    def device(self):
+        return self._sync_service.device
+
+    async def load_model(self) -> None:
+        """Async wrapper for loading the model."""
+        if self.is_ready:
+            return
+            
+        async with self._lock:
+            if self.is_ready:
+                return
+            # Run the heavy lifting in a thread
+            await asyncio.to_thread(self._sync_service.load_model)
+
+    async def generate_single_embedding(self, text: str) -> List[float]:
+        """
+        Generate a single embedding asynchronously.
+        """
+        # Ensure model is loaded first (non-blocking check mostly)
+        await self.load_model()
+        return await asyncio.to_thread(self._sync_service.generate_single_embedding, text)
+
+    async def generate_batch_embeddings(self, texts: List[str]) -> List[List[float]]:
+        """
+        Efficient batch processing asynchronously.
+        """
+        await self.load_model()
+        return await asyncio.to_thread(self._sync_service.generate_batch_embeddings, texts)

app/services/embeddings_service.py

@@ -0,0 +1,169 @@
+"""
+EmbeddingsService:
+- Loads tokenizer + model with retries (tenacity)
+- Efficiently computes embeddings for large batches by chunking into BATCH_SIZE
+- Mean-pools the last hidden states to produce embeddings
+- Handles OOM errors by falling back to micro-batch single-item processing
+"""
+
+from ..config import settings
+from ..logger import logger
+
+import math
+from typing import List, Optional
+
+import torch
+from transformers import AutoTokenizer, AutoModel
+from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
+
+
+class EmbeddingsService:
+    def __init__(self, model_name: Optional[str] = None):
+        self.model_name = model_name or settings.MODEL_NAME
+        # If DEVICE is set in env, honor it; otherwise auto-detect GPU if available
+        if settings.DEVICE:
+            self.device = torch.device(settings.DEVICE)
+        else:
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        self.model: Optional[AutoModel] = None
+        self.tokenizer: Optional[AutoTokenizer] = None
+        self.is_ready: bool = False
+
+    # Retry decorator: will attempt to load model multiple times with exponential backoff
+    @retry(
+        stop=stop_after_attempt(settings.RETRY_ATTEMPTS),
+        wait=wait_exponential(multiplier=settings.RETRY_BACKOFF_SECONDS, max=20), # max wait 20 second for next retry
+        retry=retry_if_exception_type(Exception),
+        reraise=True,
+    )
+    def load_model(self) -> None:
+        if self.is_ready:
+            return
+            
+        logger.info(f"Loading model: {self.model_name}")
+        logger.info(f"Target device: {self.device}")
+        
+        try:
+            #-------------------------------------------------
+            # Using private Hugging Face token to fetch model
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                self.model_name, 
+                use_fast=True,
+                token=settings.HUGGING_FACE_TOKEN
+            )
+            self.model = AutoModel.from_pretrained(
+                self.model_name,
+                use_safetensors=True,  
+                token=settings.HUGGING_FACE_TOKEN
+            )
+            #-------------------------------------------------
+            self.model.to(self.device)
+            self.model.eval()
+            self.is_ready = True
+            logger.info(f"Model '{self.model_name}' loaded and ready.")
+        except Exception as exc:
+            logger.error(f"Failed to load model '{self.model_name}'")
+            logger.exception(exc)
+            raise
+
+    def _pool_embeddings(self, last_hidden_state: torch.Tensor) -> torch.Tensor:
+        """Mean pooling across token dimension: (B, T, D) -> (B, D)."""
+        return last_hidden_state.mean(dim=1)
+
+    @torch.no_grad()
+    def generate_single_embedding(self, text: str) -> List[float]:
+        """
+        Generate a single embedding synchronously.
+        """
+        self.load_model()
+        if not text:
+            logger.warning("Received empty text for embedding")
+            raise ValueError("Empty text provided for embedding.")
+
+        log_text = (text[:100] + '...') if len(text) > 100 else text
+        logger.debug(f"Processing single embedding for text: '{log_text}'")
+
+        # debuging here text length by tokens
+        # token_count = len(self.tokenizer.tokenize(text))
+        # if token_count > settings.MAX_LENGTH:
+        #     logger.warning(f"Text token count {token_count} is longer than max length ({settings.MAX_LENGTH}), truncating...")
+        
+        try:
+            inputs = self.tokenizer(
+                text,
+                return_tensors="pt",
+                truncation=True,
+                padding=True,
+                max_length=settings.MAX_LENGTH,
+            ).to(self.device)
+
+            outputs = self.model(**inputs)
+            emb = self._pool_embeddings(outputs.last_hidden_state)
+
+            return emb.cpu().numpy().flatten().tolist()
+        except Exception as e:
+            logger.error("Critical error during single embedding generation")
+            logger.exception(e)
+            raise
+
+    @torch.no_grad()
+    def generate_batch_embeddings(self, texts: List[str]) -> List[List[float]]:
+        """
+        Efficient batch processing with OOM recovery (Sync).
+        """
+        self.load_model()
+        if not texts:
+            return []
+
+        batch_size = max(1, settings.BATCH_SIZE)
+        embeddings: List[List[float]] = []
+
+        total = len(texts)
+        num_chunks = math.ceil(total / batch_size)
+        logger.info(f"Processing {total} texts in {num_chunks} chunks (batch_size={batch_size})")
+        
+        for i in range(0, total, batch_size):
+            chunk = texts[i : i + batch_size]
+            try:
+                # debuging here every text
+                for index, text in enumerate(chunk):
+                    log_text = (text[:100] + '...') if len(text) > 100 else text
+                    logger.debug(f"Processing chunk-{i+1}, embedding for text-{index+1}: '{log_text}'")
+                    
+                    # debuging here text length by tokens
+                    # token_count = len(self.tokenizer.tokenize(text))
+                    # if token_count > settings.MAX_LENGTH:
+                    #      logger.warning(f"Text token count {token_count} is longer than max length ({settings.MAX_LENGTH}), truncating...")
+
+                inputs = self.tokenizer(
+                    chunk,
+                    return_tensors="pt",
+                    truncation=True,
+                    padding=True,
+                    max_length=settings.MAX_LENGTH,
+                ).to(self.device)
+
+                outputs = self.model(**inputs)
+                chunk_embs = self._pool_embeddings(outputs.last_hidden_state)
+                embeddings.extend(chunk_embs.cpu().numpy().tolist())
+            except RuntimeError as e:
+                # Basic check for OOM
+                msg = str(e).lower()
+                if "out of memory" in msg or "cuda out of memory" in msg:
+                    logger.warning(f"GPU OOM on chunk {i//batch_size + 1}. Attempting micro-batch fallback...")
+                    try:
+                        torch.cuda.empty_cache()
+                    except: pass
+                    for t in chunk:
+                        embeddings.append(self.generate_single_embedding(t))
+                else:
+                    logger.error(f"Runtime error in batch processing at chunk {i//batch_size + 1}")
+                    logger.exception(e)
+                    raise
+            except Exception as e:
+                logger.error(f"Unexpected error in batch processing at chunk {i//batch_size + 1}")
+                logger.exception(e)
+                raise
+
+        return embeddings

docker-compose.yml

@@ -0,0 +1,31 @@
+# Docker Compose file for BGE Embedding Server
+
+services:
+  embed-server:
+    build: .
+    container_name: bge_embed_server
+    environment:
+      - EMBED_API_KEY=${EMBED_API_KEY}
+      - MODEL_NAME=${MODEL_NAME}
+      - BATCH_SIZE=${BATCH_SIZE:-8}
+      - MAX_LENGTH=${MAX_LENGTH:-1024}
+      - RETRY_ATTEMPTS=${RETRY_ATTEMPTS:-3}
+      - RETRY_BACKOFF_SECONDS=${RETRY_BACKOFF_SECONDS:-2}
+      - LOG_LEVEL=${LOG_LEVEL:-INFO}
+      - PORT=${PORT}
+      - HOST=${HOST}
+      - WORKERS=${WORKERS}
+      - DEVICE=${DEVICE}
+    ports:
+      - "${PORT:-7860}:7860"
+    volumes:
+      - ./:/app
+    restart: unless-stopped
+    # If i have GPU server i simply uncomment it to use gpu
+    # deploy:
+    #   resources:
+    #     reservations:
+    #       devices:
+    #         - driver: nvidia
+    #           count: all
+    #           capabilities: [ gpu ]

prompt.py

@@ -0,0 +1,113 @@
+"""
+You are an AI Sales Assistant representing {{ORGANIZATION_NAME}}.
+Communicate in the {{organization_prefered}} style.
+
+You will receive each turn:
+- ORGANIZATION_SUMMARY: {{ORGANIZATION_SUMMARY}}
+- CHAT_HISTORY: {{CHAT_HISTORY}}
+- CURRENT USER MESSAGE: {{USER_MESSAGE}}
+- LEAD_FIELDS: {{LEAD_FIELDS_JSON}}
+- ORGANIZATION_EXTRA_RULE (optional_field): {{ORGANIZATION_EXTRA_RULE}}
+- OLD_CHAT_SUMMARY (optional_field): {{OLD_CHAT_SUMMARY}} 
+- EXTRA_DEATILS_ABOUT_THIS_LEAD (optional_field): {{EXTRA_DEATILS_ABOUT_THIS_LEAD}} 
+
+CRITICAL: Return ONLY valid JSON. No markdown, no explanation, no extra text.
+
+---
+CORE RULES
+---
+1. "user_reply" is the actual SMS sent to user. Max 4/5 sentences if possible. Keep it conversational and brief.
+2. Ask at most ONE follow-up question per reply about ONE null field in LEAD_FIELDS.
+3. Auto-fill fields from context/images when possible before asking.
+4. NEVER overwrite non-null LEAD_FIELDS unless user explicitly corrects it.
+5. If user provides info for a field, update that field in lead_fields.
+6. If calling a tool, return ONLY tool_call JSON (see TOOLS section).
+7. If user is abusive/spam, respond politely once, set "auto_reply_off": true, "escalate_to_human": false.
+8. If project is significantly larger than typical (check ORGANIZATION_SUMMARY and ORGANIZATION_EXTRA_RULE), set "escalate_to_human": true and say "Our senior team will contact you very soon."
+9. NEVER hallucinate. If you don't have info, tell user politely you'll have them contacted shortly.
+10. If user wants to book a meeting, allow it anytime - don't wait for other details. Try to get preferred date/time; if they say "anytime" put that in optional_sms.
+11. If all LEAD_FIELDS are filled (no nulls), suggest booking a meeting to discuss their project.
+
+---
+TOOLS
+---
+If using a tool, return ONLY this:
+
+{
+  "tool_call": {
+    "tool_name": "<tool_name>",
+    "parameters": { ... }
+  }
+}
+
+Available tools:
+
+search_product
+  { "query": "description of product/image" }
+  Use: user asks about products or sends product image
+
+request_quotation
+  { 
+    "message": "brief summary of what user wants quoted/ordered",
+    "details": "specific details, quantities, specs, generated demo quotation etc.",
+    "urgent": true or false
+  }
+  Use: user requests quotation, wants to order, or asks for pricing on specific items
+
+book_meeting
+  { 
+    "preferred_date": "YYYY-MM-DD or null",
+    "preferred_time": "HH:MM or null",
+    "contact_name": "string or null",
+    "optional_sms": "any additional notes like 'anytime works' or meeting purpose"
+  }
+  Use: user requests meeting/scheduling (allow anytime, even before getting other details)
+
+set_reminder
+  { 
+    "reminder_time": "YYYY-MM-DD HH:MM",
+    "note": "context note"
+  }
+  Use: user asks to be reminded later
+
+notify_human
+  { "notification_message": "what needs human attention" }
+  Use: user asks question you can't answer or needs immediate human help
+
+TOOL RESPONSE:
+Next turn after tool_call, you'll receive PREVIOUS_TOOL_CALL and TOOL_RESPONSE. Incorporate the result into conversation (max 4/5 sentences) and update lead_fields if applicable. For request_quotation and book_meeting, confirm to user they'll be contacted soon by the team.
+
+---
+OUTPUT FORMAT (when NOT calling tool)
+---
+{
+  "user_reply": "<SMS text, max 4/5 sentences>",
+  "lead_fields": {
+    "budget": <string|null>,
+    "project_duration": <string|null>,
+    "deadline": <string|null>,
+    "project_type": <string|null>
+  },
+  "escalate_to_human": <true|false>,
+  "auto_reply_off": <true|false>,
+  "tool_call": null
+}
+
+Rules:
+- lead_fields: updated state after this turn. Keep existing non-null values.
+- escalate_to_human: true only for large leads or urgent human-required situations
+- auto_reply_off: true only for spam/abuse/non-leads
+- tool_call: must be null in this format
+
+---
+BEHAVIOR
+---
+- Auto-fill fields from user message/image before asking questions
+- Ask ONE high-value question if user shows buying intent
+- Use appropriate tool when user requests product info, quotation, meeting, reminder, or asks unanswerable question
+- Allow meeting booking anytime - don't require other details first
+- If all lead_fields are filled, suggest booking a meeting to discuss their project
+- Never make up pricing, availability, discount, hallucinate, or product details
+- Be concise and natural like human in user_reply
+
+"""

requirements.lock.txt

@@ -0,0 +1,10 @@
+fastapi==0.129.0
+uvicorn[standard]==0.40.0
+transformers==5.1.0
+torch==2.10.0
+accelerate==1.12.0
+pydantic==2.12.5
+python-dotenv==1.2.1
+tenacity==9.1.4
+colorlog==6.10.1
+pydantic-settings==2.13.0

requirements.txt

@@ -0,0 +1,10 @@
+fastapi>=0.95.0
+uvicorn[standard]>=0.22.0
+transformers>=4.35.0
+torch>=2.0.0
+accelerate>=0.20.0
+pydantic>=1.10.0
+python-dotenv>=1.0.0
+tenacity>=8.2.0
+colorlog>=6.7.0
+pydantic-settings>=2.0.0

scripts/convert_safetensors.py

@@ -0,0 +1,87 @@
+#=====================================================================================================
+# This script converts a PyTorch model to SafeTensors format and uploads it to Hugging Face Hub.
+# It is used to fix the "safetensors not found" error.
+# Usage: python scripts/convert_safetensors.py
+# It need to Run only once if the "safetensors not found" error came otherwise don't run it.
+# NOTE: U need write access to the repo to upload the model.
+#=====================================================================================================
+
+import sys
+import os
+import shutil
+
+# Adjust path to find app module
+current_dir = os.path.dirname(os.path.abspath(__file__))
+parent_dir = os.path.dirname(current_dir)
+sys.path.append(parent_dir)
+
+from app.config import settings
+from transformers import AutoModel, AutoTokenizer
+from huggingface_hub import HfApi
+
+def convert():
+    model_name = settings.MODEL_NAME
+    token = settings.HUGGING_FACE_TOKEN
+    temp_dir = os.path.join(parent_dir, "temp_safe_model")
+    
+    print(f"Loading original model: {model_name}...")
+    try:
+        # Load the PyTorch version explicitly
+        model = AutoModel.from_pretrained(
+            model_name, 
+            use_safetensors=False, 
+            token=token
+        )
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_name, 
+            token=token
+        )
+    except Exception as e:
+        print(f"Failed to load original model: {e}")
+        return
+
+    print("Saving model locally with SafeTensors format...")
+    try:
+        if os.path.exists(temp_dir):
+            shutil.rmtree(temp_dir)
+        os.makedirs(temp_dir, exist_ok=True)
+            
+        model.save_pretrained(temp_dir, safe_serialization=True)
+        tokenizer.save_pretrained(temp_dir)
+        
+        # --- Verification Step ---
+        print("Verifying converted model by loading it back...")
+        try:
+            # Try to load the model from the temporary directory using SafeTensors
+            check_model = AutoModel.from_pretrained(temp_dir, use_safetensors=True)
+            print("Verification successful! Model loaded correctly from SafeTensors.")
+            # memory cleanup
+            del check_model
+        except Exception as e:
+            print(f"Verification FAILED: {e}")
+            print("Aborting upload.")
+            return
+
+        # --- Upload Step ---
+        print("Model verified. Now uploading to Hub...")
+        
+        api = HfApi(token=token)
+        api.upload_folder(
+            folder_path=temp_dir,
+            repo_id=model_name,
+            repo_type="model"
+        )
+        
+        print("Success! The model has been converted, verified, and pushed to your repository.")
+        print("The auto-conversion error should now be resolved.")
+        
+    except Exception as e:
+        print(f"An error occurred during the process: {e}")
+    finally:
+        # Cleanup
+        if os.path.exists(temp_dir):
+            shutil.rmtree(temp_dir)
+            print("Cleaned up temporary files.")
+
+if __name__ == "__main__":
+    convert()