Spaces:
Sleeping
Sleeping
| from contextlib import asynccontextmanager | |
| from fastapi import FastAPI | |
| from fastapi.responses import FileResponse | |
| from pydantic import BaseModel | |
| from translator import load_translator | |
| # Point this at YOUR weights repo on the Hugging Face Hub. | |
| WEIGHTS_REPO = "ME22B191/neural_translation_weights" | |
| CHECKPOINT_FILE = "checkpoint_baseline.pt" # <-- the exact filename you uploaded | |
| ml = {} # holds the one loaded model for the whole process lifetime | |
| async def lifespan(app: FastAPI): | |
| # Runs once at startup: download weights + vocab, assemble the model. | |
| ml["model"] = load_translator(WEIGHTS_REPO, checkpoint_file=CHECKPOINT_FILE) | |
| yield | |
| ml.clear() # runs once at shutdown | |
| app = FastAPI(title="Neural Translator", lifespan=lifespan) | |
| class TranslationRequest(BaseModel): | |
| text: str | |
| class TranslationResponse(BaseModel): | |
| translation: str | |
| def home(): | |
| return FileResponse("index.html") | |
| def health_check(): | |
| return {"status": "alive", "model_loaded": "model" in ml} | |
| def translate(request: TranslationRequest): # plain def -> runs in threadpool | |
| # infer() already does .eval() + torch.no_grad() internally. | |
| english = ml["model"].infer(request.text) | |
| return TranslationResponse(translation=english) |