Update app.py

app.py

@@ -1,57 +1,60 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import T5Tokenizer, T5ForConditionalGeneration  # Note: Using specific T5 classes
 import logging
 import os
-import torch
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
-# Initialize FastAPI
 app = FastAPI()
-
-# Get HF token
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
-# Define models
 MODELS = {
     "nidra-v1": "m1k3wn/nidra-v1",
     "nidra-v2": "m1k3wn/nidra-v2"
 }
 
-# Simple request model
 class PredictionRequest(BaseModel):
     inputs: str
     model: str = "nidra-v1"
 
-# Simple response model
 class PredictionResponse(BaseModel):
     generated_text: str
 
-@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):
     try:
-        # Load model
+        logger.info(f"Loading model: {request.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
+        # Use T5-specific classes instead of Auto classes
+        tokenizer = T5Tokenizer.from_pretrained(
+            model_path,
+            token=HF_TOKEN,
+            legacy=True  # Try with legacy mode first
+        )
+        
+        model = T5ForConditionalGeneration.from_pretrained(
+            model_path,
+            token=HF_TOKEN
+        )
+        
         full_input = "Interpret this dream: " + request.inputs
-        inputs = tokenizer(full_input, return_tensors="pt")
-        outputs = model.generate(**inputs)
+        logger.info(f"Processing: {full_input}")
+        
+        inputs = tokenizer(
+            full_input,
+            return_tensors="pt",
+            truncation=True,
+            max_length=512
+        )
+        
+        outputs = model.generate(**inputs, max_length=200)
         result = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
         return PredictionResponse(generated_text=result)
+        
     except Exception as e:
+        logger.error(f"Error: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))