Update main.py

main.py

@@ -1,17 +1,19 @@
 import os
-from typing import Dict, Any
+from typing import Dict, Any, Optional, Tuple
 from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from huggingface_hub.utils import RepositoryNotFoundError
 import torch
 
-# ---- Config -----------------------------------------------------------------
-MODEL_ID = os.getenv("MODEL_ID", "ethnmcl/checkin-lora-gpt2")  # NEW default
-BASE_TOKENIZER = os.getenv("BASE_TOKENIZER", "gpt2")           # fallback if LoRA repo has no tokenizer
-HF_TOKEN = os.getenv("HF_TOKEN")  # set if private
+# ---- Config --------------------------------------------------------------
+PREFERRED_ID = os.getenv("MODEL_ID", "ethnmcl/checkin-lora-gpt2")
+FALLBACK_IDS = ["ethnmcl/checkin-lora-gpt2", "distilgpt2"]  # last-resort keeps API alive
+BASE_TOKENIZER = os.getenv("BASE_TOKENIZER", "gpt2")
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-app = FastAPI(title="Check-in GPT-2 API", version="1.2.0")
+app = FastAPI(title="Check-in GPT-2 API", version="1.3.0")
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"],
@@ -20,51 +22,69 @@ app.add_middleware(
 device = 0 if torch.cuda.is_available() else -1
 DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# ---- Tokenizer (with fallback for adapter-only repos) ------------------------
-def load_tokenizer(repo_id: str, token: str | None):
+# ---- Helpers -------------------------------------------------------------
+def _load_tokenizer(repo_id: str) -> Tuple:
+    """Try repo tokenizer, then fallback to base tokenizer."""
     try:
-        tk = AutoTokenizer.from_pretrained(repo_id, token=token)
+        tk = AutoTokenizer.from_pretrained(repo_id, token=HF_TOKEN)
         if tk.pad_token is None:
             tk.pad_token = tk.eos_token
         return tk, repo_id, False
-    except Exception as e_model_tok:
-        # Adapter repos often don't include tokenizer files: fallback to base tokenizer
-        tk = AutoTokenizer.from_pretrained(BASE_TOKENIZER, token=token)
+    except Exception:
+        tk = AutoTokenizer.from_pretrained(BASE_TOKENIZER, token=HF_TOKEN)
         if tk.pad_token is None:
             tk.pad_token = tk.eos_token
         return tk, BASE_TOKENIZER, True
 
-tokenizer, tokenizer_source, tokenizer_fallback = load_tokenizer(MODEL_ID, HF_TOKEN)
+def _try_plain(repo_id: str):
+    return AutoModelForCausalLM.from_pretrained(
+        repo_id, token=HF_TOKEN, dtype=DTYPE,
+        device_map="auto" if torch.cuda.is_available() else None,
+    )
 
-# ---- Model (plain or PEFT LoRA) ---------------------------------------------
-_merged = False
-try:
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_ID,
-        token=HF_TOKEN,
-        dtype=DTYPE,
+def _try_peft(repo_id: str):
+    from peft import AutoPeftModelForCausalLM
+    m = AutoPeftModelForCausalLM.from_pretrained(
+        repo_id, token=HF_TOKEN, dtype=DTYPE,
         device_map="auto" if torch.cuda.is_available() else None,
     )
-except Exception as e_plain:
-    # Try PEFT (adapter) path
+    # Merge if available; ok if not
     try:
-        from peft import AutoPeftModelForCausalLM
-        model = AutoPeftModelForCausalLM.from_pretrained(
-            MODEL_ID,
-            token=HF_TOKEN,
-            dtype=DTYPE,
-            device_map="auto" if torch.cuda.is_available() else None,
-        )
+        m = m.merge_and_unload()
+        merged = True
+    except Exception:
+        merged = False
+    return m, merged
+
+def load_model_any(repo_id: str):
+    """Try plain, then PEFT; raise if both fail."""
+    try:
+        m = _try_plain(repo_id)
+        return m, False
+    except Exception as e_plain:
         try:
-            model = model.merge_and_unload()
-            _merged = True
-        except Exception:
-            _merged = False
-    except Exception as e_peft:
-        raise RuntimeError(
-            f"Failed to load model '{MODEL_ID}'. "
-            f"Plain load error: {e_plain}\nPEFT load error: {e_peft}"
-        )
+            m, merged = _try_peft(repo_id)
+            return m, merged
+        except Exception as e_peft:
+            raise RuntimeError(f"load failed for {repo_id} | plain: {e_plain} | peft: {e_peft}")
+
+# ---- Boot: try MODEL_ID first, then fallbacks ----------------------------
+errors = {}
+chosen_id: Optional[str] = None
+merged_lora = False
+
+trial_ids = [PREFERRED_ID] + [i for i in FALLBACK_IDS if i != PREFERRED_ID]
+for rid in trial_ids:
+    try:
+        tokenizer, tokenizer_source, tokenizer_fallback_used = _load_tokenizer(rid)
+        model, merged_lora = load_model_any(rid)
+        chosen_id = rid
+        break
+    except Exception as e:
+        errors[rid] = str(e)
+
+if chosen_id is None:
+    raise RuntimeError(f"All model loads failed. Errors: {errors}")
 
 pipe = pipeline(
     "text-generation",
@@ -73,7 +93,7 @@ pipe = pipeline(
     device=device,
 )
 
-# ---- Prompting ---------------------------------------------------------------
+# ---- Prompting -----------------------------------------------------------
 PREFIX = "INPUT: "
 SUFFIX = "\nOUTPUT:"
 def make_prompt(user_input: str) -> str:
@@ -98,11 +118,13 @@ class GenerateResponse(BaseModel):
 def root():
     return {
         "message": "Check-in GPT-2 API. POST /generate",
-        "model": MODEL_ID,
+        "model_chosen": chosen_id,
         "device": "cuda" if device == 0 else "cpu",
-        "merged_lora": _merged,
+        "merged_lora": merged_lora,
         "tokenizer_source": tokenizer_source,
-        "tokenizer_fallback_used": tokenizer_fallback,
+        "tokenizer_fallback_used": tokenizer_fallback_used,
+        "attempt_errors": errors,
+        "env_MODEL_ID": PREFERRED_ID,
     }
 
 @app.get("/health")
@@ -131,3 +153,4 @@ def generate(req: GenerateRequest):
         return GenerateResponse(output=output, prompt=prompt, parameters=req.model_dump())
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
+