Update main.py

main.py

@@ -6,34 +6,72 @@ from pydantic import BaseModel, Field
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 import torch
 
-MODEL_ID = os.getenv("MODEL_ID", "ethnmcl/checkin-gpt2")
-HF_TOKEN = os.getenv("HF_TOKEN")  # set in Space Secrets if the model repo is private
-
-app = FastAPI(title="Check-in GPT-2 API", version="1.0.0")
+# === 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
 
+app = FastAPI(title="Check-in GPT-2 API", version="1.1.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
 
-# ✅ use token= (not use_auth_token) and rely on HF_HOME=/data/huggingface
+# === Load tokenizer ===
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
 if tokenizer.pad_token is None:
     tokenizer.pad_token = tokenizer.eos_token
-model = AutoModelForCausalLM.from_pretrained(MODEL_ID, token=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
+_merged = False
+try:
+    _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,
+        device_map="auto" if torch.cuda.is_available() else None,
+    )
+except Exception as e_plain:
+    # Fall back to PEFT path
+    try:
+        from peft import AutoPeftModelForCausalLM
+        _model = AutoPeftModelForCausalLM.from_pretrained(
+            MODEL_ID,
+            token=HF_TOKEN,
+            dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+            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()
+            _merged = True
+        except Exception:
+            # If merge not available, we still can run with adapters active
+            _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}"
+        )
+
+# Build pipeline
 pipe = pipeline(
     "text-generation",
-    model=model,
+    model=_model,
     tokenizer=tokenizer,
-    device=device
+    device=device,
 )
 
+# Prompt shape (keep if you rely on INPUT/OUTPUT markers; otherwise switch to 'Check-in: ')
 PREFIX = "INPUT: "
 SUFFIX = "\nOUTPUT:"
-
 def make_prompt(user_input: str) -> str:
     return f"{PREFIX}{user_input}{SUFFIX}"
 
@@ -54,7 +92,12 @@ class GenerateResponse(BaseModel):
 
 @app.get("/")
 def root():
-    return {"message": "Check-in GPT-2 API. POST /generate", "model": MODEL_ID, "device": "cuda" if device == 0 else "cpu"}
+    return {
+        "message": "Check-in GPT-2 API. POST /generate",
+        "model": MODEL_ID,
+        "device": "cuda" if device == 0 else "cpu",
+        "merged_lora": _merged,
+    }
 
 @app.get("/health")
 def health():
@@ -75,7 +118,7 @@ def generate(req: GenerateRequest):
             num_return_sequences=req.num_return_sequences,
             pad_token_id=tokenizer.eos_token_id,
             eos_token_id=tokenizer.eos_token_id,
-            return_full_text=True
+            return_full_text=True,
         )
         text = gen[0]["generated_text"]
         output = text.split("OUTPUT:", 1)[-1].strip()