Create main.py

main.py

@@ -0,0 +1,89 @@
+import os
+from typing import Dict, Any
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+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 repo is private
+PORT = int(os.getenv("PORT", "7860"))
+
+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)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+model = AutoModelForCausalLM.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN)
+
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    device=device
+)
+
+PREFIX = "INPUT: "
+SUFFIX = "\nOUTPUT:"
+
+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")
+    max_new_tokens: int = 180
+    temperature: float = 0.7
+    top_p: float = 0.95
+    top_k: int = 50
+    repetition_penalty: float = 1.05
+    do_sample: bool = True
+    num_return_sequences: int = 1
+
+class GenerateResponse(BaseModel):
+    output: str
+    prompt: str
+    parameters: Dict[str, Any]
+
+@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"
+    }
+
+@app.get("/health")
+def health():
+    return {"status": "ok"}
+
+@app.post("/generate", response_model=GenerateResponse)
+def generate(req: GenerateRequest):
+    try:
+        prompt = make_prompt(req.input)
+        gen = pipe(
+            prompt,
+            max_new_tokens=req.max_new_tokens,
+            temperature=req.temperature,
+            top_p=req.top_p,
+            top_k=req.top_k,
+            repetition_penalty=req.repetition_penalty,
+            do_sample=req.do_sample,
+            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
+        )
+        text = gen[0]["generated_text"]
+        output = text.split("OUTPUT:", 1)[-1].strip()
+        return GenerateResponse(output=output, prompt=prompt, parameters=req.model_dump())
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))