Spaces:

maxxcarl
/

emb1024

Sleeping

gcharanteja

ch6

ad469c9 6 months ago

5.64 kB

from fastapi import FastAPI, Header, HTTPException, Body
from sentence_transformers import SentenceTransformer
import uvicorn
import os
from pathlib import Path
import chromadb
from chromadb.config import Settings
from chromadb.server.fastapi import FastAPI as ChromaFastAPI
import torch
import logging
import asyncio
from concurrent.futures import ThreadPoolExecutor
from typing import List

# --- Logging Setup ---
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(process)d - %(levelname)s - %(message)s",
)
logger = logging.getLogger("HarrierService")

# --- CPU Optimization ---
# Crucial for multi-worker setups: prevents workers from fighting over the same CPU cores
torch.set_num_threads(1)

app = FastAPI(title="Harrier OSS 0.6B Embedder Production Service")

# --- Configuration ---
MODEL_NAME = "microsoft/harrier-oss-v1-0.6b"
API_KEY_REQUIRED = os.getenv("API_KEY", "Azure123")
QUERY_INSTRUCTION = "Instruct: Retrieve relevant passages that answer the query\nQuery: "
CHROMA_HOST = os.getenv("CHROMA_HOST", "0.0.0.0")
CHROMA_PERSIST_DIRECTORY = os.getenv("CHROMA_PERSIST_DIRECTORY") or "/data/chroma_data"

# Force CPU since it's a CPU-only server
device = "cpu"
logger.info(f"[*] Starting service worker on device: {device}")

# Internal thread pool for each worker to handle the handoff
executor = ThreadPoolExecutor(max_workers=1)

# Model loaded globally per worker process
model = None


def ensure_bucket_path(path: str) -> str:
    bucket_path = Path(path)
    try:
        bucket_path.mkdir(parents=True, exist_ok=True)
    except OSError as exc:
        raise RuntimeError(
            f"Storage path is not writable: {bucket_path}. "
            "This server must run on a Hugging Face Space with a mounted bucket."
        ) from exc
    if not os.access(bucket_path, os.W_OK):
        raise RuntimeError(
            f"Storage path is not writable: {bucket_path}. "
            "This server must run on a Hugging Face Space with a mounted bucket."
        )
    return str(bucket_path)


chroma_persist_directory = ensure_bucket_path(CHROMA_PERSIST_DIRECTORY)
chroma_settings = Settings(persist_directory=chroma_persist_directory)
chroma_server = ChromaFastAPI(chroma_settings)
app.mount("/chroma", chroma_server.app())

def write_bucket_probe(path: str) -> None:
    probe_path = Path(path) / "test.md"
    probe_path.write_text(
        "Chroma bucket probe: if you can read this file, the Space can write to /data.\n",
        encoding="utf-8",
    )


def seed_chroma_data(path: str) -> None:
    client = chromadb.PersistentClient(path=path)
    collection = client.get_or_create_collection(name="knowledge_base")
    if collection.count() > 0:
        logger.info("[*] Chroma already has data; skipping seed.")
        return

    documents = [
        "Chroma is a lightweight, open-source vector database built for AI.",
        "Python is a high-level programming language used extensively in data science.",
        "The celestial body closest to Earth is the Moon.",
    ]
    metadatas = [
        {"category": "tech", "source": "docs"},
        {"category": "tech", "source": "wiki"},
        {"category": "science", "source": "space-facts"},
    ]
    ids = ["doc1", "doc2", "doc3"]
    collection.add(documents=documents, metadatas=metadatas, ids=ids)
    logger.info("[+] Seeded Chroma with sample documents.")


@app.on_event("startup")
def load_model():
    global model
    write_bucket_probe(chroma_persist_directory)
    seed_chroma_data(chroma_persist_directory)
    logger.info(f"[*] Loading Harrier OSS 0.6B model: {MODEL_NAME}...")
    try:
        model = SentenceTransformer(MODEL_NAME, trust_remote_code=True, device=device)
        logger.info("[+] Model loaded successfully into worker.")
    except Exception as e:
        logger.error(f"[-] Failed to load model: {e}")
        raise e


def compute_embeddings(processed_input):
    """Heavy computation logic."""
    return model.encode(processed_input, show_progress_bar=False).tolist()


@app.post("/embed")
async def get_embeddings(
    input: List[str] = Body(...),
    is_query: bool = Body(False),
    x_api_key: str = Header(None),
):
    # 1. API Key Validation
    if x_api_key != API_KEY_REQUIRED:
        raise HTTPException(status_code=403, detail="Invalid API Key")

    # 2. Input Validation
    if not input:
        raise HTTPException(status_code=400, detail="Input list cannot be empty")

    # 3. Apply Harrier-specific instruction
    processed_input = [f"{QUERY_INSTRUCTION}{text}" for text in input] if is_query else input

    # 4. Offload to ThreadPool
    loop = asyncio.get_event_loop()
    try:
        logger.info(f"[*] Processing {len(input)} items (is_query={is_query})")
        embeddings = await loop.run_in_executor(executor, compute_embeddings, processed_input)
        return {"embeddings": embeddings}
    except Exception as e:
        logger.error(f"Error during embedding: {e}")
        raise HTTPException(status_code=500, detail="Internal embedding error")


@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "model": MODEL_NAME,
        "device": device,
        "worker_pid": os.getpid(),
    }


if __name__ == "__main__":
    # --- Deployment Config for 16GB RAM / CPU ---
    # With a 0.6B model (~2.5GB RAM), 4 workers = ~10GB.
    # This leaves 6GB for OS and overhead, which is safe and handles concurrency well.
    uvicorn.run(
        "embedder_service:app",
        host="0.0.0.0",
        workers=4,  # This creates 4 independent processes
        access_log=False,  # Disable for slightly better performance
    )

	from fastapi import FastAPI, Header, HTTPException, Body
	from sentence_transformers import SentenceTransformer
	import uvicorn
	import os
	from pathlib import Path
	import chromadb
	from chromadb.config import Settings
	from chromadb.server.fastapi import FastAPI as ChromaFastAPI
	import torch
	import logging
	import asyncio
	from concurrent.futures import ThreadPoolExecutor
	from typing import List

	# --- Logging Setup ---
	logging.basicConfig(
	level=logging.INFO,
	format="%(asctime)s - %(process)d - %(levelname)s - %(message)s",
	)
	logger = logging.getLogger("HarrierService")

	# --- CPU Optimization ---
	# Crucial for multi-worker setups: prevents workers from fighting over the same CPU cores
	torch.set_num_threads(1)

	app = FastAPI(title="Harrier OSS 0.6B Embedder Production Service")

	# --- Configuration ---
	MODEL_NAME = "microsoft/harrier-oss-v1-0.6b"
	API_KEY_REQUIRED = os.getenv("API_KEY", "Azure123")
	QUERY_INSTRUCTION = "Instruct: Retrieve relevant passages that answer the query\nQuery: "
	CHROMA_HOST = os.getenv("CHROMA_HOST", "0.0.0.0")
	CHROMA_PERSIST_DIRECTORY = os.getenv("CHROMA_PERSIST_DIRECTORY") or "/data/chroma_data"

	# Force CPU since it's a CPU-only server
	device = "cpu"
	logger.info(f"[*] Starting service worker on device: {device}")

	# Internal thread pool for each worker to handle the handoff
	executor = ThreadPoolExecutor(max_workers=1)

	# Model loaded globally per worker process
	model = None


	def ensure_bucket_path(path: str) -> str:
	bucket_path = Path(path)
	try:
	bucket_path.mkdir(parents=True, exist_ok=True)
	except OSError as exc:
	raise RuntimeError(
	f"Storage path is not writable: {bucket_path}. "
	"This server must run on a Hugging Face Space with a mounted bucket."
	) from exc
	if not os.access(bucket_path, os.W_OK):
	raise RuntimeError(
	f"Storage path is not writable: {bucket_path}. "
	"This server must run on a Hugging Face Space with a mounted bucket."
	)
	return str(bucket_path)


	chroma_persist_directory = ensure_bucket_path(CHROMA_PERSIST_DIRECTORY)
	chroma_settings = Settings(persist_directory=chroma_persist_directory)
	chroma_server = ChromaFastAPI(chroma_settings)
	app.mount("/chroma", chroma_server.app())

	def write_bucket_probe(path: str) -> None:
	probe_path = Path(path) / "test.md"
	probe_path.write_text(
	"Chroma bucket probe: if you can read this file, the Space can write to /data.\n",
	encoding="utf-8",
	)


	def seed_chroma_data(path: str) -> None:
	client = chromadb.PersistentClient(path=path)
	collection = client.get_or_create_collection(name="knowledge_base")
	if collection.count() > 0:
	logger.info("[*] Chroma already has data; skipping seed.")
	return

	documents = [
	"Chroma is a lightweight, open-source vector database built for AI.",
	"Python is a high-level programming language used extensively in data science.",
	"The celestial body closest to Earth is the Moon.",
	]
	metadatas = [
	{"category": "tech", "source": "docs"},
	{"category": "tech", "source": "wiki"},
	{"category": "science", "source": "space-facts"},
	]
	ids = ["doc1", "doc2", "doc3"]
	collection.add(documents=documents, metadatas=metadatas, ids=ids)
	logger.info("[+] Seeded Chroma with sample documents.")


	@app.on_event("startup")
	def load_model():
	global model
	write_bucket_probe(chroma_persist_directory)
	seed_chroma_data(chroma_persist_directory)
	logger.info(f"[*] Loading Harrier OSS 0.6B model: {MODEL_NAME}...")
	try:
	model = SentenceTransformer(MODEL_NAME, trust_remote_code=True, device=device)
	logger.info("[+] Model loaded successfully into worker.")
	except Exception as e:
	logger.error(f"[-] Failed to load model: {e}")
	raise e


	def compute_embeddings(processed_input):
	"""Heavy computation logic."""
	return model.encode(processed_input, show_progress_bar=False).tolist()


	@app.post("/embed")
	async def get_embeddings(
	input: List[str] = Body(...),
	is_query: bool = Body(False),
	x_api_key: str = Header(None),
	):
	# 1. API Key Validation
	if x_api_key != API_KEY_REQUIRED:
	raise HTTPException(status_code=403, detail="Invalid API Key")

	# 2. Input Validation
	if not input:
	raise HTTPException(status_code=400, detail="Input list cannot be empty")

	# 3. Apply Harrier-specific instruction
	processed_input = [f"{QUERY_INSTRUCTION}{text}" for text in input] if is_query else input

	# 4. Offload to ThreadPool
	loop = asyncio.get_event_loop()
	try:
	logger.info(f"[*] Processing {len(input)} items (is_query={is_query})")
	embeddings = await loop.run_in_executor(executor, compute_embeddings, processed_input)
	return {"embeddings": embeddings}
	except Exception as e:
	logger.error(f"Error during embedding: {e}")
	raise HTTPException(status_code=500, detail="Internal embedding error")


	@app.get("/health")
	async def health():
	return {
	"status": "healthy",
	"model": MODEL_NAME,
	"device": device,
	"worker_pid": os.getpid(),
	}


	if __name__ == "__main__":
	# --- Deployment Config for 16GB RAM / CPU ---
	# With a 0.6B model (~2.5GB RAM), 4 workers = ~10GB.
	# This leaves 6GB for OS and overhead, which is safe and handles concurrency well.
	uvicorn.run(
	"embedder_service:app",
	host="0.0.0.0",
	workers=4, # This creates 4 independent processes
	access_log=False, # Disable for slightly better performance
	)