Update app.py

app.py

@@ -1,18 +1,18 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import T5Tokenizer, T5ForConditionalGeneration  # Changed to specific classes
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 import logging
 import os
 import sys
 
-@app.get("/version")
-async def version():
-    return {"python_version": sys.version}
+# Initialize FastAPI first
+app = FastAPI()
 
+# Set up logging
 logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger(__name__)
 
-app = FastAPI()
+# Get HF token
 HF_TOKEN = os.environ.get("HF_TOKEN")
 
 MODELS = {
@@ -27,6 +27,14 @@ class PredictionRequest(BaseModel):
 class PredictionResponse(BaseModel):
     generated_text: str
 
+@app.get("/version")
+async def version():
+    return {"python_version": sys.version}
+
+@app.get("/health")
+async def health():
+    return {"status": "healthy"}
+
 @app.post("/predict", response_model=PredictionResponse)
 async def predict(request: PredictionRequest):
     try:
@@ -38,7 +46,7 @@ async def predict(request: PredictionRequest):
         tokenizer = T5Tokenizer.from_pretrained(
             model_path,
             token=HF_TOKEN,
-            local_files_only=False,  # Force download if needed
+            local_files_only=False,
             return_special_tokens_mask=True
         )
         logger.info("Tokenizer loaded successfully")
@@ -47,7 +55,7 @@ async def predict(request: PredictionRequest):
         model = T5ForConditionalGeneration.from_pretrained(
             model_path,
             token=HF_TOKEN,
-            local_files_only=False  # Force download if needed
+            local_files_only=False
         )
         logger.info("Model loaded successfully")
         
@@ -76,11 +84,6 @@ async def predict(request: PredictionRequest):
     except Exception as e:
         logger.error(f"Error: {str(e)}")
         logger.error(f"Error type: {type(e)}")
-        # Log the full traceback
         import traceback
         logger.error(f"Traceback: {traceback.format_exc()}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.get("/health")
-async def health():
-    return {"status": "healthy"}
+        raise HTTPException(status_code=500, detail=str(e))