Update app.py

simplifies for debugging

app.py

@@ -1,8 +1,7 @@
 from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel, validator
+from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import logging
-from typing import Optional, Dict, Any
 import os
 import torch
 
@@ -10,124 +9,49 @@ import torch
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-app = FastAPI(title="Dream Interpretation API")
+# Initialize FastAPI
+app = FastAPI()
 
-# Get HF token from environment variable
+# Get HF token
 HF_TOKEN = os.environ.get("HF_TOKEN")
-if not HF_TOKEN:
-    raise ValueError("HF_TOKEN environment variable must be set")
 
-# Define the model names
+# Define models
 MODELS = {
     "nidra-v1": "m1k3wn/nidra-v1",
     "nidra-v2": "m1k3wn/nidra-v2"
 }
 
-# Cache for loaded models
-loaded_models = {}
-loaded_tokenizers = {}
-
-# Pydantic models for request/response validation
+# Simple request model
 class PredictionRequest(BaseModel):
     inputs: str
     model: str = "nidra-v1"
-    parameters: Optional[Dict[str, Any]] = {}
-
-    @validator('inputs')
-    def validate_inputs(cls, v):
-        if not isinstance(v, str):
-            raise ValueError('inputs must be a string')
-        if not v.strip():
-            raise ValueError('inputs cannot be empty')
-        return v.strip()
-
-    @validator('model')
-    def validate_model(cls, v):
-        if v not in MODELS:
-            raise ValueError(f'model must be one of: {", ".join(MODELS.keys())}')
-        return v
 
+# Simple response model
 class PredictionResponse(BaseModel):
     generated_text: str
 
-def load_model(model_name: str):
-    """Load model and tokenizer on demand"""
-    if model_name not in loaded_models:
-        logger.info(f"Loading {model_name}...")
-        try:
-            model_path = MODELS[model_name]
-            
-            logger.info("Loading tokenizer...")
-            tokenizer = AutoTokenizer.from_pretrained(
-                model_path, 
-                token=HF_TOKEN,
-                use_fast=False
-            )
-            
-            logger.info("Loading model...")
-            model = AutoModelForSeq2SeqLM.from_pretrained(
-                model_path,
-                token=HF_TOKEN,
-                torch_dtype=torch.float32,
-            )
-            
-            model = model.cpu()
-            
-            loaded_models[model_name] = model
-            loaded_tokenizers[model_name] = tokenizer
-            logger.info(f"Successfully loaded {model_name}")
-        except Exception as e:
-            logger.error(f"Error loading {model_name}: {str(e)}")
-            raise
-    return loaded_tokenizers[model_name], loaded_models[model_name]
+@app.get("/")
+async def root():
+    return {"message": "Dream Interpretation API", "status": "running"}
+
+@app.get("/health")
+async def health():
+    return {"status": "healthy"}
 
 @app.post("/predict", response_model=PredictionResponse)
 async def predict(request: PredictionRequest):
-    """Make a prediction using the specified model"""
     try:
-        # Load model on demand
-        tokenizer, model = load_model(request.model)
-
-        # Log the input for debugging
-        logger.info(f"Processing input: {request.inputs}")
-
-        # Prepend the shared prefix
+        # Load model
+        model_path = MODELS[request.model]
+        tokenizer = AutoTokenizer.from_pretrained(model_path, token=HF_TOKEN)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_path, token=HF_TOKEN)
+        
+        # Process input
         full_input = "Interpret this dream: " + request.inputs
-        logger.info(f"Full input: {full_input}")
-
-        try:
-            # Tokenize
-            tokenizer_output = tokenizer(
-                full_input, 
-                return_tensors="pt",
-                padding=True,
-                truncation=True,
-                max_length=512
-            )
-            logger.info("Tokenization successful")
-            
-            input_ids = tokenizer_output.input_ids
-            
-            # Generate
-            outputs = model.generate(
-                input_ids,
-                max_length=200,
-                num_return_sequences=1,
-                no_repeat_ngram_size=2,
-                **request.parameters
-            )
-            logger.info("Generation successful")
-            
-            # Decode
-            decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)
-            logger.info(f"Decoded output: {decoded}")
-
-        except Exception as e:
-            logger.error(f"Error in model prediction pipeline: {str(e)}")
-            raise HTTPException(status_code=500, detail=f"Model prediction failed: {str(e)}")
-
-        return PredictionResponse(generated_text=decoded)
-
+        inputs = tokenizer(full_input, return_tensors="pt")
+        outputs = model.generate(**inputs)
+        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        return PredictionResponse(generated_text=result)
     except Exception as e:
-        logger.error(f"Error in prediction: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))