app.py

import torch
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from transformers import T5Tokenizer, T5ForConditionalGeneration, GenerationConfig
from typing import Optional, Dict, Any, ClassVar
import logging
import os
import sys
import traceback
from functools import lru_cache
import gc
import asyncio
from fastapi import BackgroundTasks
import psutil

# Initialize FastAPI
app = FastAPI()

# Debugging logs
logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Get HF token
HF_TOKEN = os.environ.get("HF_TOKEN")
if not HF_TOKEN:
    logger.warning("No HF_TOKEN found in environment variables")

MODELS = {
    "nidra-v1": "m1k3wn/nidra-v1",
    "nidra-v2": "m1k3wn/nidra-v2"
}

DEFAULT_GENERATION_CONFIGS = {
    "nidra-v1": {
        "max_length": 300,
        "min_length": 150,
        "num_beams": 8,
        "temperature": 0.55,
        "do_sample": True,
        "top_p": 0.95,
        "repetition_penalty": 4.5,
        "no_repeat_ngram_size": 4,
        "early_stopping": True,
        "length_penalty": 1.2,
    },
    "nidra-v2": {
        "max_length": 300,
        "min_length": 150,
        "num_beams": 8,
        "temperature": 0.4,
        "do_sample": True,
        "top_p": 0.95,
        "repetition_penalty": 3.5,
        "no_repeat_ngram_size": 4,
        "early_stopping": True,
        "length_penalty": 1.2,
    }
}

class ModelManager:
    _instances: ClassVar[Dict[str, tuple]] = {}

    @classmethod
    async def get_model_and_tokenizer(cls, model_name: str):
        if model_name not in cls._instances:
            try:
                model_path = MODELS[model_name]
                logger.debug(f"Loading tokenizer and model from {model_path}")
                
                tokenizer = T5Tokenizer.from_pretrained(
                    model_path,
                    token=HF_TOKEN,
                    use_fast=True
                )
                
                model = T5ForConditionalGeneration.from_pretrained(
                    model_path,
                    token=HF_TOKEN,
                    torch_dtype=torch.float32,
                    low_cpu_mem_usage=True,
                    device_map='auto'
                )
                
                model.eval()
                torch.set_num_threads(6)  # Number of CPUs used
                cls._instances[model_name] = (model, tokenizer)
                
            except Exception as e:
                logger.error(f"Error loading {model_name}: {str(e)}")
                raise
                
        return cls._instances[model_name]

class PredictionRequest(BaseModel):
    inputs: str
    model: str = "nidra-v1"
    parameters: Optional[Dict[str, Any]] = None

class PredictionResponse(BaseModel):
    generated_text: str
    selected_model: str  # Changed from model_used to avoid namespace conflict

# Memory debug endpoint
@app.get("/debug/memory")
async def memory_usage():
    process = psutil.Process()
    memory_info = process.memory_info()
    return {
        "memory_used_mb": memory_info.rss / 1024 / 1024,
        "memory_percent": process.memory_percent(),
        "cpu_percent": process.cpu_percent()
    }

# Version check
@app.get("/version")
async def version():
    return {
        "python_version": sys.version,
        "models_available": list(MODELS.keys())
    }

# Healthcheck endpoint
@app.get("/health")
async def health():
    try:
        logger.debug("Health check started")
        logger.debug(f"HF_TOKEN present: {bool(HF_TOKEN)}")
        logger.debug(f"Available models: {MODELS}")
        
        result = await ModelManager.get_model_and_tokenizer("nidra-v1")
        logger.debug("Model and tokenizer loaded successfully")
        
        return {
            "status": "healthy",
            "loaded_models": list(ModelManager._instances.keys())
        }
    except Exception as e:
        error_msg = f"Health check failed: {str(e)}\n{traceback.format_exc()}"
        logger.error(error_msg)
        return {
            "status": "unhealthy",
            "error": str(e)
        }

@app.post("/predict", response_model=PredictionResponse)
async def predict(request: PredictionRequest, background_tasks: BackgroundTasks):
    try:
        if request.model not in MODELS:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid model. Available models: {list(MODELS.keys())}"
            )

        model, tokenizer = await ModelManager.get_model_and_tokenizer(request.model)
        
        # Add immediate cleanup of memory before generation
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

        generation_params = DEFAULT_GENERATION_CONFIGS[request.model].copy()

        try:
            model_generation_config = model.generation_config
            generation_params.update({
                k: v for k, v in model_generation_config.to_dict().items()
                if v is not None
            })
        except Exception as config_load_error:
            logger.warning(f"Using default generation config: {config_load_error}")

        if request.parameters:
            generation_params.update(request.parameters)

        logger.debug(f"Final generation parameters: {generation_params}")

        full_input = "Interpret this dream: " + request.inputs
        inputs = tokenizer(
            full_input,
            return_tensors="pt",
            truncation=True,
            max_length=512,
            padding=True,
            return_attention_mask=True
        )

        async def generate():
            try:
                return model.generate(
                    **inputs,
                    **{k: v for k, v in generation_params.items() if k in [
                        'max_length', 'min_length', 'do_sample', 'temperature',
                        'top_p', 'top_k', 'num_beams', 'no_repeat_ngram_size',
                        'repetition_penalty', 'early_stopping'
                    ]}
                )
            finally:
                # Ensure cleanup happens even if generation fails
                gc.collect()
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
            
        with torch.inference_mode():
            outputs = await asyncio.wait_for(generate(), timeout=45.0)  # Reduced timeout
            
        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
        background_tasks.add_task(cleanup_memory)
        
        return PredictionResponse(
            generated_text=result,
            selected_model=request.model
        )

    except asyncio.TimeoutError:
        logger.error("Generation timed out")
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        raise HTTPException(status_code=504, detail="Generation timed out")
    except Exception as e:
        error_msg = f"Error during prediction: {str(e)}\n{traceback.format_exc()}"
        logger.error(error_msg)
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        raise HTTPException(status_code=500, detail=error_msg)

def cleanup_memory():
    try:
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        # Force Python garbage collection
        gc.collect(generation=2)
        
    except Exception as e:
        logger.error(f"Error in cleanup: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)