fix: lifespan handler, dtype= for transformers 5.x, remove UI, update README

README.md

@@ -8,64 +8,88 @@ pinned: false
 app_port: 7860
 ---
 
-# LFM2-700M — FastAPI Inference Space
+# LFM2-700M — FastAPI Inference API
 
-Runs [gsstec/LFM2-700M](https://huggingface.co/gsstec/LFM2-700M) behind a lightweight **FastAPI** server.
+Runs [gsstec/LFM2-700M](https://huggingface.co/gsstec/LFM2-700M) behind a lightweight **FastAPI** server on CPU.
 
 ## Endpoints
 
 | Method | Path | Description |
 |--------|------|-------------|
-| `GET`  | `/`  | Interactive HTML docs page |
-| `GET`  | `/health` | Liveness check |
-| `POST` | `/chat`   | Generate a response |
+| `GET`  | `/health` | Liveness + readiness check |
+| `POST` | `/chat`   | Multi-turn chat (supports streaming) |
+| `POST` | `/ask`    | Single-turn question |
 
-## POST /chat — request body
+---
+
+## GET /health
+
+```json
+{
+  "status": "ok",
+  "model": "gsstec/LFM2-700M",
+  "ready": true,
+  "dtype": "torch.float32",
+  "cpu_threads": 16,
+  "authenticated": true
+}
+```
+
+---
+
+## POST /chat
 
 ```json
 {
   "messages": [
-    {"role": "user", "content": "What is C. elegans?"}
+    {"role": "system", "content": "You are a helpful e-commerce assistant."},
+    {"role": "user",   "content": "How do I reduce cart abandonment?"}
   ],
-  "max_new_tokens": 512,
+  "max_new_tokens": 256,
   "temperature": 0.3,
   "min_p": 0.15,
   "repetition_penalty": 1.05,
-  "stream": false
+  "stream": true
 }
 ```
 
-Set `"stream": true` to receive a streaming `text/plain` response via chunked transfer.
+Set `"stream": true` to receive a `text/plain` chunked streaming response.  
+Set `"stream": false` to receive:
+
+```json
+{ "response": "Here are some tips to reduce cart abandonment..." }
+```
+
+---
 
-## POST /chat — response body (non-streaming)
+## POST /ask
 
 ```json
 {
-  "response": "C. elegans, also known as Caenorhabditis elegans, is a small ..."
+  "question": "What is a good return policy for an online store?",
+  "max_new_tokens": 128,
+  "temperature": 0.3
 }
 ```
 
+Response:
+
+```json
+{ "response": "A good return policy should be..." }
+```
+
+---
+
 ## Local development
 
 ```bash
 pip install -r requirements.txt
-uvicorn app:app --reload --port 7860
+HF_TOKEN=hf_... python app.py
 ```
 
 ## Docker
 
 ```bash
 docker build -t lfm2-700m .
-docker run -p 7860:7860 lfm2-700m
+docker run -e HF_TOKEN=hf_... -p 7860:7860 lfm2-700m
 ```
-
-## Deploy to HuggingFace Spaces
-
-1. Create a new Space at https://huggingface.co/new-space → choose **Docker** SDK
-2. Push all files to the Space repo:
-   ```bash
-   git remote add space https://huggingface.co/spaces/<your-org>/<space-name>
-   git push space main
-   ```
-
-> **GPU tip:** Upgrade the Space hardware to a GPU tier and uncomment `attn_implementation="flash_attention_2"` in `app.py` for faster inference.

app.py

@@ -3,6 +3,7 @@ import json
 import torch
 import logging
 import traceback
+from contextlib import asynccontextmanager
 
 # ---------------------------------------------------------------------------
 # Environment must be configured BEFORE importing transformers/huggingface_hub
@@ -10,15 +11,13 @@ import traceback
 os.environ.setdefault("HF_HOME", "/app/model_cache")
 os.environ.setdefault("TRANSFORMERS_CACHE", "/app/model_cache")
 
-# Propagate HF_TOKEN so authenticated requests are made automatically
 _hf_token = os.getenv("HF_TOKEN")
 if _hf_token:
     os.environ["HUGGING_FACE_HUB_TOKEN"] = _hf_token
-    os.environ["HF_TOKEN"] = _hf_token
 
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.exceptions import RequestValidationError
-from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
+from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
 from typing import List
 from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
@@ -30,32 +29,32 @@ logging.getLogger("uvicorn.access").setLevel(logging.INFO)
 
 MODEL_ID = os.getenv("MODEL_ID", "gsstec/LFM2-700M")
 
-# Max out CPU threads for faster matrix ops
 _num_threads = int(os.getenv("NUM_THREADS", os.cpu_count() or 4))
 torch.set_num_threads(_num_threads)
 torch.set_num_interop_threads(max(1, _num_threads // 2))
 
-# BF16 is not reliably supported on all CPUs — fall back to float32
-_dtype = torch.bfloat16 if torch.backends.cpu.get_cpu_capability() >= "avx512" else torch.float32
-
-# ---------------------------------------------------------------------------
-# App — model is loaded in the startup event so FastAPI controls timing
-# and any exception is visible in logs rather than killing the process silently
-# ---------------------------------------------------------------------------
-app = FastAPI(title="LFM2-700M Inference API")
+# BF16 needs AVX-512; HF CPU Spaces are typically AVX2 → use float32
+_dtype = (
+    torch.bfloat16
+    if torch.backends.cpu.get_cpu_capability() >= "avx512"
+    else torch.float32
+)
 
 model = None
 tokenizer = None
 _model_ready = False
 
 
-@app.on_event("startup")
-async def load_model():
+# ---------------------------------------------------------------------------
+# Lifespan — loads model after uvicorn binds the port
+# ---------------------------------------------------------------------------
+@asynccontextmanager
+async def lifespan(app: FastAPI):
     global model, tokenizer, _model_ready
     _token_kwarg = {"token": _hf_token} if _hf_token else {}
 
     print(f"MODEL: {MODEL_ID}")
-    print(f"dtype: {_dtype}  threads: {_num_threads}  authenticated: {bool(_hf_token)}")
+    print(f"dtype={_dtype}  threads={_num_threads}  authenticated={bool(_hf_token)}")
 
     try:
         print("Loading tokenizer...")
@@ -72,9 +71,8 @@ async def load_model():
         print("Loading model...")
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
-            torch_dtype=_dtype,
+            dtype=_dtype,
             low_cpu_mem_usage=True,
-            # attn_implementation="flash_attention_2",  # uncomment on GPU
             **_token_kwarg,
         )
         model.eval()
@@ -83,12 +81,16 @@ async def load_model():
 
     except Exception:
         traceback.print_exc()
-        # Don't crash uvicorn — /health will report not ready
+
+    yield   # server is running
 
 
 # ---------------------------------------------------------------------------
-# Validation error handler
+# App
 # ---------------------------------------------------------------------------
+app = FastAPI(title="LFM2-700M Inference API", lifespan=lifespan)
+
+
 @app.exception_handler(RequestValidationError)
 async def validation_exception_handler(request: Request, exc: RequestValidationError):
     body = await request.body()
@@ -97,7 +99,6 @@ async def validation_exception_handler(request: Request, exc: RequestValidationE
          for k, v in err.items()}
         for err in exc.errors()
     ]
-    logger.error(f"422 {request.method} {request.url} | {safe_errors}")
     return JSONResponse(
         status_code=422,
         content={"detail": safe_errors, "body_received": body.decode(errors="replace")},
@@ -108,7 +109,7 @@ async def validation_exception_handler(request: Request, exc: RequestValidationE
 # Schemas
 # ---------------------------------------------------------------------------
 class Message(BaseModel):
-    role: str       # "user" | "assistant" | "system"
+    role: str
     content: str
 
 
@@ -118,7 +119,7 @@ class ChatRequest(BaseModel):
     temperature: float = 0.3
     min_p: float = 0.15
     repetition_penalty: float = 1.05
-    stream: bool = True          # default to streaming on CPU
+    stream: bool = True
 
 
 class ChatResponse(BaseModel):
@@ -134,7 +135,7 @@ class AskRequest(BaseModel):
 
 
 # ---------------------------------------------------------------------------
-# Shared helpers
+# Helpers
 # ---------------------------------------------------------------------------
 def _check_ready():
     if not _model_ready:
@@ -172,16 +173,6 @@ def _generate(input_ids: torch.Tensor, req, streamer=None) -> torch.Tensor:
 # ---------------------------------------------------------------------------
 # Routes
 # ---------------------------------------------------------------------------
-_UI_HTML = open(os.path.join(os.path.dirname(__file__), "static", "index.html"), encoding="utf-8").read()
-
-
-@app.get("/", response_class=HTMLResponse)
-@app.get("/ui", response_class=HTMLResponse)
-def ui():
-    """Chat UI — served at both / and /ui"""
-    return _UI_HTML
-
-
 @app.get("/health")
 def health():
     return {