app.py

import os

cache_dir = os.path.join(os.getcwd(), ".cache")
try:
    os.makedirs(cache_dir, exist_ok=True)
except PermissionError:
    cache_dir = "/tmp/.cache"
    os.makedirs(cache_dir, exist_ok=True)

os.environ['TRANSFORMERS_CACHE'] = cache_dir
os.environ['HF_HOME'] = cache_dir


from fastapi import FastAPI, Header, HTTPException, Body
from transformers import AutoTokenizer

# 從環境變數讀取 secret
EXPECTED_API_KEY = os.environ.get("apikey")

# 初始化 tokenizer
tokenizer = AutoTokenizer.from_pretrained("ckiplab/bert-base-chinese-ws")

app = FastAPI(title="CKIP Word Segmentation API")

# ✅ GET 方法：健康檢查
@app.get("/check")
def health_check():
    return {"status": "ok", "message": "API is running"}

# ✅ POST 方法：斷詞
@app.post("/tokenize")
async def tokenize(
    text: str = Body(..., embed=True),
    x_api_key: str = Header(None)
):
    # 驗證 API Key 設定是否存在
    if not EXPECTED_API_KEY:
        raise HTTPException(status_code=500, detail="Server missing API_KEY config")
    # 驗證 API Key 值是否正確
    if x_api_key != EXPECTED_API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API Key")

    tokens = tokenizer.tokenize(text)
    return {"tokens": tokens}