app.py

import io, json, logging, os, time, uuid, asyncio
from collections import defaultdict, deque

import torch
from fastapi import FastAPI, HTTPException, Depends, Request
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.responses import JSONResponse
from huggingface_hub import HfApi
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
from peft import PeftModel
from pydantic import BaseModel

# ---------- Config ----------
BASE_MODEL_ID = "wolethereader/STORM-OS-MT-3B"
REVERSE_ADAPTER_ID = "wolethereader/STORM-OS-MT-3B-REVERSE"
HF_TOKEN = os.environ.get("HF_TOKEN")  # required: both model repos are private
ORG_NAME = "wolethereader"
EN = "eng_Latn"
LANG_CODES = {"yo": "yor_Latn", "ha": "hau_Latn", "ig": "ibo_Latn", "pcm": "pcm_Latn"}
VALID_LANGS = set(LANG_CODES.keys())
MAX_TEXT_CHARS = 2000

app = FastAPI(title="STORM-OS MT API")

def log_event(event, **fields):
    print(json.dumps({"event": event, "ts": time.time(), **fields}))

# ---------- Auth: HF token AND must belong to the org ----------
security = HTTPBearer()
hf_api = HfApi()
_token_cache = {}
TOKEN_CACHE_TTL = 300

EXTERNAL_ACCESS_TOKEN = os.environ.get("EXTERNAL_ACCESS_TOKEN")

async def verify_org_token(creds: HTTPAuthorizationCredentials = Depends(security)):
    token = creds.credentials

    if EXTERNAL_ACCESS_TOKEN and token == EXTERNAL_ACCESS_TOKEN:
        log_event("auth_external_token_used")
        return "external-collaborator"

    now = time.time()
    cached = _token_cache.get(token)
    if cached and cached[1] > now:
        return cached[0]
    try:
        info = hf_api.whoami(token=token)
    except Exception:
        log_event("auth_failed_invalid_token")
        raise HTTPException(status_code=401, detail="Invalid or expired Hugging Face token")
    username = info.get("name", "unknown")
    user_orgs = [o.get("name") for o in info.get("orgs", [])]
    if ORG_NAME not in user_orgs:
        log_event("auth_failed_not_org_member", user=username, orgs=user_orgs)
        raise HTTPException(status_code=403, detail=f"Token does not belong to a member of '{ORG_NAME}'")
    _token_cache[token] = (username, now + TOKEN_CACHE_TTL)
    return username

# ---------- Rate limiting ----------
_rate_state = defaultdict(deque)
RATE_LIMIT_PER_MIN = 30

def check_rate_limit(username: str):
    now = time.time()
    q = _rate_state[username]
    while q and q[0] < now - 60:
        q.popleft()
    if len(q) >= RATE_LIMIT_PER_MIN:
        raise HTTPException(status_code=429, detail="Rate limit exceeded, try again shortly")
    q.append(now)

# ---------- Model ----------
tokenizer = None
model = None  # PeftModel with the reverse adapter loaded but not necessarily active

@app.on_event("startup")
async def startup():
    global tokenizer, model
    log_event("loading_tokenizer_with_pcm_token", repo=REVERSE_ADAPTER_ID)
    tokenizer = AutoTokenizer.from_pretrained(REVERSE_ADAPTER_ID, token=HF_TOKEN)  # has pcm_Latn added

    log_event("loading_base_model", repo=BASE_MODEL_ID)
    base_model = AutoModelForSeq2SeqLM.from_pretrained(BASE_MODEL_ID, torch_dtype=torch.bfloat16, token=HF_TOKEN)
    base_model.resize_token_embeddings(len(tokenizer))

    device = "cuda" if torch.cuda.is_available() else "cpu"
    base_model.to(device)

    log_event("attaching_reverse_adapter", repo=REVERSE_ADAPTER_ID)
    model = PeftModel.from_pretrained(base_model, REVERSE_ADAPTER_ID, adapter_name="reverse", token=HF_TOKEN)
    model.eval()
    log_event("model_loaded_ok", device=device)

class TranslateRequest(BaseModel):
    text: str
    direction: str  # "forward" (local -> English) or "reverse" (English -> local)
    lang: str        # the local language code, yo/ha/ig/pcm, regardless of direction
    max_new_tokens: int = 128

@app.get("/")
def root():
    return {"status": "ok", "languages": sorted(VALID_LANGS), "directions": ["forward", "reverse"], "engine": BASE_MODEL_ID}

@app.get("/health")
def health():
    return {"status": "ok" if model is not None else "loading"}

@app.post("/translate")
async def translate(req: TranslateRequest, username: str = Depends(verify_org_token)):
    check_rate_limit(username)

    if req.lang not in VALID_LANGS:
        raise HTTPException(status_code=400, detail=f"lang must be one of {sorted(VALID_LANGS)}")
    if req.direction not in ("forward", "reverse"):
        raise HTTPException(status_code=400, detail="direction must be 'forward' or 'reverse'")
    if not req.text or not req.text.strip():
        raise HTTPException(status_code=400, detail="text must not be empty")
    if len(req.text) > MAX_TEXT_CHARS:
        raise HTTPException(status_code=400, detail=f"text exceeds {MAX_TEXT_CHARS} character limit")

    request_id = str(uuid.uuid4())
    start = time.time()

    if req.direction == "forward":
        src_lang, tgt_lang = LANG_CODES[req.lang], EN
        context = model.disable_adapter()
    else:
        src_lang, tgt_lang = EN, LANG_CODES[req.lang]
        model.set_adapter("reverse")
        context = None

    tokenizer.src_lang = src_lang
    inputs = tokenizer(req.text, return_tensors="pt", truncation=True, max_length=128).to(model.device)
    tgt_id = tokenizer.convert_tokens_to_ids(tgt_lang)

    with torch.no_grad():
        if context is not None:
            with context:
                out = model.generate(**inputs, forced_bos_token_id=tgt_id, max_new_tokens=req.max_new_tokens, max_length=None)
        else:
            out = model.generate(**inputs, forced_bos_token_id=tgt_id, max_new_tokens=req.max_new_tokens, max_length=None)

    translated = tokenizer.decode(out[0], skip_special_tokens=True)
    elapsed_s = round(time.time() - start, 2)

    log_event("translate_ok", request_id=request_id, user=username,
              direction=req.direction, lang=req.lang, elapsed_s=elapsed_s)

    return {
        "request_id": request_id,
        "direction": req.direction,
        "lang": req.lang,
        "translated_text": translated,
        "elapsed_s": elapsed_s,
    }

@app.exception_handler(HTTPException)
async def http_exception_handler(request: Request, exc: HTTPException):
    log_event("request_error", path=str(request.url.path), status_code=exc.status_code, detail=exc.detail)
    return JSONResponse(status_code=exc.status_code, content={"detail": exc.detail})