Update main.py

main.py

@@ -6,54 +6,59 @@ from pydantic import BaseModel, Field
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 
-# === Config ===
-MODEL_ID = os.getenv("MODEL_ID", "ethnmcl/checkin-lora-gpt2")
-HF_TOKEN = os.getenv("HF_TOKEN")  # if the repo is private, set this in Secrets
+# ---- 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
 
-app = FastAPI(title="Check-in GPT-2 API", version="1.1.0")
+app = FastAPI(title="Check-in GPT-2 API", version="1.2.0")
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"],
 )
 
-# Choose device: GPU index 0 if available else CPU
 device = 0 if torch.cuda.is_available() else -1
+DTYPE = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# === Load tokenizer ===
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
-if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token
+# ---- Tokenizer (with fallback for adapter-only repos) ------------------------
+def load_tokenizer(repo_id: str, token: str | None):
+    try:
+        tk = AutoTokenizer.from_pretrained(repo_id, token=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)
+        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)
 
-# === Load model (supports plain CausalLM repos AND PEFT LoRA adapters) ===
-# Strategy:
-#   1) Try plain AutoModelForCausalLM
-#   2) If that fails (likely LoRA-only repo), try PEFT AutoPeftModelForCausalLM and merge
-_model = None
+# ---- Model (plain or PEFT LoRA) ---------------------------------------------
 _merged = False
 try:
-    _model = AutoModelForCausalLM.from_pretrained(
+    model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
         token=HF_TOKEN,
-        # Use 'dtype' not deprecated 'torch_dtype'
-        dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+        dtype=DTYPE,
         device_map="auto" if torch.cuda.is_available() else None,
     )
 except Exception as e_plain:
-    # Fall back to PEFT path
+    # Try PEFT (adapter) path
     try:
         from peft import AutoPeftModelForCausalLM
-        _model = AutoPeftModelForCausalLM.from_pretrained(
+        model = AutoPeftModelForCausalLM.from_pretrained(
             MODEL_ID,
             token=HF_TOKEN,
-            dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            dtype=DTYPE,
             device_map="auto" if torch.cuda.is_available() else None,
         )
-        # Merge LoRA into base weights so inference behaves like a standard CausalLM
         try:
-            _model = _model.merge_and_unload()
+            model = model.merge_and_unload()
             _merged = True
         except Exception:
-            # If merge not available, we still can run with adapters active
             _merged = False
     except Exception as e_peft:
         raise RuntimeError(
@@ -61,15 +66,14 @@ except Exception as e_plain:
             f"Plain load error: {e_plain}\nPEFT load error: {e_peft}"
         )
 
-# Build pipeline
 pipe = pipeline(
     "text-generation",
-    model=_model,
+    model=model,
     tokenizer=tokenizer,
     device=device,
 )
 
-# Prompt shape (keep if you rely on INPUT/OUTPUT markers; otherwise switch to 'Check-in: ')
+# ---- Prompting ---------------------------------------------------------------
 PREFIX = "INPUT: "
 SUFFIX = "\nOUTPUT:"
 def make_prompt(user_input: str) -> str:
@@ -97,6 +101,8 @@ def root():
         "model": MODEL_ID,
         "device": "cuda" if device == 0 else "cpu",
         "merged_lora": _merged,
+        "tokenizer_source": tokenizer_source,
+        "tokenizer_fallback_used": tokenizer_fallback,
     }
 
 @app.get("/health")
@@ -125,4 +131,3 @@ 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))
-