File size: 2,538 Bytes
9b3bab5 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 | """
Minimal API wrapper around CorvidaeAviary so the model does something visible
when deployed as an HF Docker Space.
Endpoints:
GET / -- basic info
GET /health -- readiness probe
POST /forward -- run a forward pass on random/user-supplied token ids
"""
import os
from typing import List, Optional
import torch
from fastapi import FastAPI
from pydantic import BaseModel
from corvid_aviary import CorvidaeAviary, SPECIALIST_ORDER
VOCAB_SIZE = 200
EMBEDDING_DIM = 64
MAX_SEQ_LEN = 32
app = FastAPI(title="Corvidae Aviary")
_model: Optional[CorvidaeAviary] = None
def get_model() -> CorvidaeAviary:
global _model
if _model is None:
torch.manual_seed(0)
_model = CorvidaeAviary(
num_embeddings=VOCAB_SIZE, embedding_dim=EMBEDDING_DIM, max_seq_len=MAX_SEQ_LEN,
nhead=4, dim_feedforward=128, memory_size=32, memory_word_size=16,
num_read_heads=2, num_classes=VOCAB_SIZE, num_rook_experts=4, num_speakers=3,
num_tracked_agents=3, raven_buffer_size=4, crow_num_tools=4, crow_max_steps=3,
hippocampal_num_slots=16, hippocampal_coord_dim=3, statistical_memory_size=16,
identity_capacity=8, num_surface_contexts=4,
)
_model.eval()
return _model
class ForwardRequest(BaseModel):
token_ids: Optional[List[int]] = None # if omitted, random tokens are used
seq_len: int = 16
@app.get("/")
def root():
return {
"name": "Corvidae Aviary",
"specialists": SPECIALIST_ORDER,
"endpoints": ["/health", "/forward"],
}
@app.get("/health")
def health():
return {"status": "ok"}
@app.post("/forward")
def forward(req: ForwardRequest):
model = get_model()
if req.token_ids:
ids = req.token_ids[: MAX_SEQ_LEN]
x = torch.tensor([ids], dtype=torch.long)
else:
x = torch.randint(0, VOCAB_SIZE, (1, min(req.seq_len, MAX_SEQ_LEN)))
model.reset_memory(batch_size=x.size(0))
with torch.no_grad():
logits, species_info = model(x, return_species_info=True)
pred_ids = logits.argmax(dim=-1).squeeze(0).tolist()
route_weights = species_info["route_weights"].mean(dim=(0, 1)).tolist()
return {
"input_ids": x.squeeze(0).tolist(),
"predicted_next_ids": pred_ids,
"route_weights_by_specialist": dict(zip(SPECIALIST_ORDER, route_weights)),
}
if __name__ == "__main__":
import uvicorn
uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", 7860))) |