Update main.py

main.py

@@ -8,6 +8,7 @@ from functools import lru_cache
 import pandas as pd
 from fastapi import FastAPI, File, UploadFile, HTTPException, Body
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse, RedirectResponse
 from pydantic import BaseModel, Field
 
 from huggingface_hub import login, snapshot_download
@@ -15,7 +16,7 @@ import joblib
 import xgboost as xgb
 import numpy as np
 import torch
-from transformers import pipeline
+from transformers import AutoTokenizer, pipeline
 
 # -------- Config --------
 HF_TOKEN = (
@@ -34,7 +35,7 @@ app = FastAPI(
         f"Models:\n- {XGB_REPO}\n- {GPT2_REPO}\n"
         "Use /docs for interactive testing."
     ),
-    version="1.0.1",
+    version="1.1.0",
 )
 
 # CORS (allow all; tighten for production)
@@ -46,7 +47,42 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# -------- Model loading --------
+# -------- Convenience root & health --------
+@app.get("/", include_in_schema=False)
+def root():
+    return JSONResponse(
+        {
+            "ok": True,
+            "message": "Entrepreneur Readiness API is running.",
+            "docs": "/docs",
+            "endpoints": ["/health", "/readiness", "/score", "/score_csv", "/summarize", "/score_and_summarize"],
+        }
+    )
+
+# Liveness-only (no model load)
+@app.get("/health", include_in_schema=False)
+def health():
+    return JSONResponse({"ok": True, "status": "live", "docs": "/docs"})
+
+# Readiness (loads models)
+@app.get("/readiness")
+def readiness():
+    try:
+        _load_models()
+        return {"ok": True, "status": "ready"}
+    except Exception as e:
+        return JSONResponse({"ok": False, "status": "not_ready", "error": str(e)}, status_code=503)
+
+# Optional warm-up to trigger downloads/caching
+@app.post("/warmup", include_in_schema=False)
+def warmup():
+    try:
+        _load_models()
+        return {"ok": True, "warmed": True}
+    except Exception as e:
+        return JSONResponse({"ok": False, "error": str(e)}, status_code=500)
+
+# -------- Model loading helpers --------
 def _find_file(dirpath: str, candidates: T.Sequence[str], fallback_exts: T.Sequence[str] = ()) -> str:
     for name in candidates:
         p = os.path.join(dirpath, name)
@@ -63,7 +99,7 @@ def _download_artifacts() -> T.Tuple[str, str]:
         try:
             login(token=HF_TOKEN, add_to_git_credential=True)
         except Exception:
-            # Public models still download
+            # Continue if public
             pass
     xgb_local = snapshot_download(repo_id=XGB_REPO, token=HF_TOKEN, revision=None)
     gpt_local = snapshot_download(repo_id=GPT2_REPO, token=HF_TOKEN, revision=None)
@@ -73,7 +109,7 @@ def _download_artifacts() -> T.Tuple[str, str]:
 def _load_models():
     xgb_dir, gpt_dir = _download_artifacts()
 
-    # Preprocessor
+    # ---- Preprocessor ----
     preproc_path = _find_file(
         xgb_dir,
         candidates=[
@@ -86,7 +122,7 @@ def _load_models():
     )
     preprocessor = joblib.load(preproc_path)
 
-    # Booster
+    # ---- XGB booster ----
     booster_path = _find_file(
         xgb_dir,
         candidates=[
@@ -102,15 +138,25 @@ def _load_models():
     booster = xgb.Booster()
     booster.load_model(booster_path)
 
-    # GPT-2 text generation
+    # ---- GPT-2 text generation: robust tokenizer selection ----
     device = 0 if torch.cuda.is_available() else -1
+    try:
+        tok = AutoTokenizer.from_pretrained(gpt_dir, use_fast=True, trust_remote_code=False)
+    except Exception:
+        # Fallback for "ModelWrapper" tokenizer.json parse errors
+        tok = AutoTokenizer.from_pretrained(gpt_dir, use_fast=False, trust_remote_code=False)
+    # Ensure a pad token (map to eos if absent) to avoid generation warnings/errors
+    if tok.pad_token is None and tok.eos_token is not None:
+        tok.pad_token = tok.eos_token
+
     text_gen = pipeline(
         "text-generation",
         model=gpt_dir,
-        tokenizer=gpt_dir,
+        tokenizer=tok,
         device=device,
         trust_remote_code=False,
     )
+
     return preprocessor, booster, text_gen, xgb_dir
 
 # -------- Utils --------
@@ -142,16 +188,17 @@ def _format_prompt(inputs: dict, score: float) -> str:
     )
 
 def _summarize(inputs: dict, score: float, text_gen) -> str:
-    out = text_gen(
+    generated = text_gen(
         _format_prompt(inputs, score),
         max_new_tokens=120,
         do_sample=True,
         temperature=0.7,
         top_p=0.9,
         num_return_sequences=1,
-        eos_token_id=None,
+        eos_token_id=text_gen.tokenizer.eos_token_id,
+        pad_token_id=text_gen.tokenizer.eos_token_id,
     )[0]["generated_text"]
-    return out.split("Summary:", 1)[-1].strip() if "Summary:" in out else out.strip()
+    return generated.split("Summary:", 1)[-1].strip() if "Summary:" in generated else generated.strip()
 
 # -------- Schemas (Pydantic v2) --------
 class ScoreRequest(BaseModel):
@@ -178,14 +225,6 @@ class ScoreAndSummarizeResponse(BaseModel):
     results: T.List[ScoreAndSummarizeItem]
 
 # -------- Endpoints --------
-@app.get("/health")
-def health():
-    try:
-        _load_models()
-        return {"ok": True}
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
 @app.post("/score", response_model=ScoreResponse)
 def score_json(req: ScoreRequest = Body(...)):
     preprocessor, booster, _, _ = _load_models()