Update app.py

debugging generation configs

app.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import T5Tokenizer, T5ForConditionalGeneration
+from transformers import T5Tokenizer, T5ForConditionalGeneration, GenerationConfig
 import logging
 import os
 import sys
@@ -20,9 +20,37 @@ MODELS = {
     "nidra-v2": "m1k3wn/nidra-v2"
 }
 
+# Define default generation configurations for each model
+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 PredictionRequest(BaseModel):
     inputs: str
     model: str = "nidra-v1"
+    parameters: Optional[Dict[str, Any]] = None  # Allow custom parameters
 
 class PredictionResponse(BaseModel):
     generated_text: str
@@ -38,6 +66,10 @@ async def health():
 @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: {request.model}")
+            
         logger.info(f"Loading model: {request.model}")
         model_path = MODELS[request.model]
         
@@ -58,6 +90,27 @@ async def predict(request: PredictionRequest):
             local_files_only=False
         )
         logger.info("Model loaded successfully")
+
+        # Priority: 1. Request parameters, 2. Model's saved generation_config, 3. Default configs
+        generation_params = DEFAULT_GENERATION_CONFIGS[request.model].copy()
+
+        # Try to load model's saved generation config
+        try:
+            model_generation_config = GenerationConfig.from_pretrained(model_path)
+            # Convert to dict to merge with default configs
+            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"Could not load model's generation config: {config_load_error}")
+
+        # Override with request-specific parameters if provided
+        if request.parameters:
+            generation_params.update(request.parameters)
+
+        logger.info(f"Final Generation Parameters: {generation_params}")
+
         
         full_input = "Interpret this dream: " + request.inputs
         logger.info(f"Processing input: {full_input}")
@@ -73,7 +126,16 @@ async def predict(request: PredictionRequest):
         logger.info("Input tokenized successfully")
         
         logger.info("Generating output...")
-        outputs = model.generate(**inputs, max_length=200)
+        
+       # Generate with final parameters
+        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'
+            ]}
+        )
         logger.info("Output generated successfully")
         
         result = tokenizer.decode(outputs[0], skip_special_tokens=True)