Update main.py

main.py

@@ -7,23 +7,22 @@ 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 repo is private
-PORT = int(os.getenv("PORT", "7860"))
+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")
 
-# Allow your frontend(s)
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"],
 )
 
-# Load model once
 device = 0 if torch.cuda.is_available() else -1
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN)
+
+# ✅ use token= (not use_auth_token) and rely on HF_HOME=/data/huggingface
+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, use_auth_token=HF_TOKEN)
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, token=HF_TOKEN)
 
 pipe = pipeline(
     "text-generation",
@@ -39,7 +38,7 @@ def make_prompt(user_input: str) -> str:
     return f"{PREFIX}{user_input}{SUFFIX}"
 
 class GenerateRequest(BaseModel):
-    input: str = Field(..., min_length=1, description="Short check-in line to expand")
+    input: str = Field(..., min_length=1)
     max_new_tokens: int = 180
     temperature: float = 0.7
     top_p: float = 0.95
@@ -55,11 +54,7 @@ class GenerateResponse(BaseModel):
 
 @app.get("/")
 def root():
-    return {
-        "message": "Check-in GPT-2 API (POST /generate). Swagger: /docs",
-        "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"}
 
 @app.get("/health")
 def health():
@@ -87,3 +82,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))
+