handler.py

import runpod
import torch
from sentence_transformers import SentenceTransformer
from transformers import BitsAndBytesConfig
from typing import Dict, Union
import logging
import os
from config import Config

# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Global model instance
model = None

def _get_compute_dtype(device: str) -> torch.dtype:
    """Pick the lightest dtype supported for the current device."""
    if device == "cuda":
        try:
            major, _ = torch.cuda.get_device_capability()
            if major >= 8:
                return torch.bfloat16
        except Exception:
            pass
        return torch.float16
    return torch.float32

def _build_quantization_config(device: str, compute_dtype: torch.dtype):
    """Return BitsAndBytesConfig and resolved quantization label."""
    if device != "cuda":
        return None

    if Config.QUANTIZATION in {"4bit", "int4", "nf4"}:
        return BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_compute_dtype=torch.float16 if compute_dtype == torch.float32 else compute_dtype,
            bnb_4bit_use_double_quant=True,
            bnb_4bit_quant_type="nf4"
        )
    elif Config.QUANTIZATION in {"8bit", "int8"}:
        return BitsAndBytesConfig(load_in_8bit=True)
    
    return None

def init_model():
    """Initialize the model"""
    global model
    
    logger.info(f"Loading model {Config.MODEL_NAME}...")
    device = "cuda" if torch.cuda.is_available() else "cpu"
    compute_dtype = _get_compute_dtype(device)
    quantization_config = _build_quantization_config(device, compute_dtype)
    
    model_kwargs = {"torch_dtype": compute_dtype}
    if quantization_config:
        model_kwargs["quantization_config"] = quantization_config
        model_kwargs["device_map"] = "auto"
        
    try:
        model = SentenceTransformer(
            Config.MODEL_NAME,
            device=device,
            trust_remote_code=True,
            model_kwargs=model_kwargs
        )
        
        if hasattr(model, "max_seq_length") and model.max_seq_length:
            model.max_seq_length = min(model.max_seq_length, Config.MAX_SEQ_LENGTH)
            
        logger.info("Model loaded successfully")
        
    except Exception as e:
        logger.error(f"Failed to load model: {e}")
        raise

def handler(job):
    """RunPod handler function"""
    job_input = job["input"]
    
    # Extract input
    input_data = job_input.get("input")
    if not input_data:
        return {"error": "No input provided"}
        
    if isinstance(input_data, str):
        input_texts = [input_data]
    else:
        input_texts = input_data
        
    # Optional parameters
    dimensions = job_input.get("dimensions")
    
    # Generate embeddings
    try:
        embeddings = model.encode(
            input_texts,
            batch_size=Config.EMBED_BATCH_SIZE,
            convert_to_numpy=True,
            show_progress_bar=False
        )
        
        # Convert to list
        if getattr(embeddings, "ndim", 1) == 1:
            embeddings_list = [embeddings.tolist()]
        else:
            embeddings_list = embeddings.tolist()
            
        # Handle dimensions
        if dimensions:
            base_dim = len(embeddings_list[0])
            trim_dim = min(dimensions, base_dim)
            embeddings_list = [v[:trim_dim] for v in embeddings_list]
            
        # Format response (OpenAI compatible-ish)
        data = []
        total_tokens = 0
        
        for i, embedding in enumerate(embeddings_list):
            data.append({
                "object": "embedding",
                "embedding": embedding,
                "index": i
            })
            # Rough token count
            total_tokens += len(input_texts[i].split())
            
        return {
            "object": "list",
            "data": data,
            "model": Config.MODEL_NAME,
            "usage": {
                "prompt_tokens": total_tokens,
                "total_tokens": total_tokens
            }
        }
        
    except Exception as e:
        logger.error(f"Inference error: {e}")
        return {"error": str(e)}

# Initialize model at startup
init_model()

# Start RunPod
runpod.serverless.start({"handler": handler})