Update app.py

app.py

@@ -7,14 +7,58 @@ import logging
 import os
 import sys
 import traceback
-import psutil
 from functools import lru_cache
 
-[... rest of your existing code until ModelManager class ...]
+# Initialize FastAPI
+app = FastAPI()
+
+# Set up logging with more detailed formatting
+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
     def get_model_and_tokenizer(cls, model_name: str):
         if model_name not in cls._instances:
@@ -29,13 +73,26 @@ class ModelManager:
                 )
                 
                 logger.info(f"Loading model {model_name}")
-                model = T5ForConditionalGeneration.from_pretrained(
-                    model_path,
-                    token=HF_TOKEN,
-                    local_files_only=False,
-                    low_cpu_mem_usage=True,
-                    torch_dtype=torch.float32
-                ).cpu()
+                # Check if accelerate is available
+                try:
+                    import accelerate
+                    logger.info("Accelerate package found, using device_map='auto'")
+                    model = T5ForConditionalGeneration.from_pretrained(
+                        model_path,
+                        token=HF_TOKEN,
+                        local_files_only=False,
+                        device_map="auto",
+                        low_cpu_mem_usage=True,
+                        torch_dtype=torch.float32
+                    )
+                except ImportError:
+                    logger.warning("Accelerate package not found, falling back to CPU")
+                    model = T5ForConditionalGeneration.from_pretrained(
+                        model_path,
+                        token=HF_TOKEN,
+                        local_files_only=False
+                    )
+                    model = model.cpu()
                 
                 cls._instances[model_name] = (model, tokenizer)
                 logger.info(f"Successfully loaded {model_name}")
@@ -45,9 +102,17 @@ class ModelManager:
                     status_code=500,
                     detail=f"Failed to load model {model_name}: {str(e)}"
                 )
+        
         return cls._instances[model_name]
 
-[... rest of your existing code until before @app.get("/version") ...]
+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
 
 @app.get("/debug/memory")
 async def memory_usage():
@@ -59,4 +124,94 @@ async def memory_usage():
         "cpu_percent": process.cpu_percent()
     }
 
-[... rest of your existing code ...]
+@app.get("/version")
+async def version():
+    return {
+        "python_version": sys.version,
+        "models_available": list(MODELS.keys())
+    }
+
+@app.get("/health")
+async def health():
+    # More comprehensive health check
+    try:
+        # Try to load at least one model to verify functionality
+        ModelManager.get_model_and_tokenizer("nidra-v1")
+        return {
+            "status": "healthy",
+            "loaded_models": list(ModelManager._instances.keys())
+        }
+    except Exception as e:
+        logger.error(f"Health check failed: {str(e)}")
+        return {
+            "status": "unhealthy",
+            "error": str(e)
+        }
+
+@app.post("/predict", response_model=PredictionResponse)
+async def predict(request: PredictionRequest):
+    try:
+        # Validate model
+        if request.model not in MODELS:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Invalid model. Available models: {list(MODELS.keys())}"
+            )
+
+        # Get cached model and tokenizer
+        model, tokenizer = ModelManager.get_model_and_tokenizer(request.model)
+
+        # Get generation parameters
+        generation_params = DEFAULT_GENERATION_CONFIGS[request.model].copy()
+
+        # Try to load model's saved generation config
+        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}")
+
+        # Override with request-specific parameters
+        if request.parameters:
+            generation_params.update(request.parameters)
+
+        logger.debug(f"Final generation parameters: {generation_params}")
+
+        # Prepare input
+        full_input = "Interpret this dream: " + request.inputs
+        inputs = tokenizer(
+            full_input,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512,
+            padding=True
+        ).to(model.device)  # Ensure inputs are on same device as model
+
+        # Generate
+        outputs = 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'
+            ]}
+        )
+
+        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        return PredictionResponse(
+            generated_text=result,
+            selected_model=request.model
+        )
+
+    except Exception as e:
+        error_msg = f"Error during prediction: {str(e)}\n{traceback.format_exc()}"
+        logger.error(error_msg)
+        raise HTTPException(status_code=500, detail=error_msg)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)