Update app.py

app.py

@@ -1,8 +1,9 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import T5Tokenizer, T5ForConditionalGeneration  # Note: Using specific T5 classes
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 import logging
 import os
+import json
 
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
@@ -15,6 +16,17 @@ MODELS = {
     "nidra-v2": "m1k3wn/nidra-v2"
 }
 
+# Define the tokenizer configuration explicitly
+TOKENIZER_CONFIG = {
+    "model_max_length": 512,
+    "clean_up_tokenization_spaces": False,
+    "tokenizer_class": "T5Tokenizer",
+    "pad_token": "<pad>",
+    "eos_token": "</s>",
+    "unk_token": "<unk>",
+    "extra_ids": 100
+}
+
 class PredictionRequest(BaseModel):
     inputs: str
     model: str = "nidra-v1"
@@ -28,33 +40,56 @@ async def predict(request: PredictionRequest):
         logger.info(f"Loading model: {request.model}")
         model_path = MODELS[request.model]
         
-        # Use T5-specific classes instead of Auto classes
+        # Initialize tokenizer with explicit config
         tokenizer = T5Tokenizer.from_pretrained(
             model_path,
             token=HF_TOKEN,
-            legacy=True  # Try with legacy mode first
+            model_max_length=TOKENIZER_CONFIG["model_max_length"],
+            clean_up_tokenization_spaces=TOKENIZER_CONFIG["clean_up_tokenization_spaces"],
+            pad_token=TOKENIZER_CONFIG["pad_token"],
+            eos_token=TOKENIZER_CONFIG["eos_token"],
+            unk_token=TOKENIZER_CONFIG["unk_token"],
+            extra_ids=TOKENIZER_CONFIG["extra_ids"],
+            use_fast=True  # Try forcing the fast tokenizer
         )
         
         model = T5ForConditionalGeneration.from_pretrained(
             model_path,
-            token=HF_TOKEN
+            token=HF_TOKEN,
+            torch_dtype="auto"
         )
         
         full_input = "Interpret this dream: " + request.inputs
         logger.info(f"Processing: {full_input}")
         
+        # Add explicit encoding parameters
         inputs = tokenizer(
             full_input,
             return_tensors="pt",
             truncation=True,
-            max_length=512
+            max_length=512,
+            padding=True,
+            add_special_tokens=True
+        )
+        
+        outputs = model.generate(
+            **inputs,
+            max_length=200,
+            num_beams=4,
+            no_repeat_ngram_size=2,
+            length_penalty=1.0
         )
         
-        outputs = model.generate(**inputs, max_length=200)
         result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        logger.info(f"Generated result: {result}")
         
         return PredictionResponse(generated_text=result)
         
     except Exception as e:
         logger.error(f"Error: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=str(e))
+
+# Add health check endpoint
+@app.get("/health")
+async def health():
+    return {"status": "healthy"}