from __future__ import annotations import json from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer import torch from tiny_transformer.train import load_checkpoint HTML = """ Tiny Transformer Playground

Tiny Transformer Playground


  

  


"""


def serve_playground(checkpoint: str, host: str, port: int, device: str) -> None:
    model, tokenizer = load_checkpoint(checkpoint, device=device)

    class Handler(BaseHTTPRequestHandler):
        def do_GET(self) -> None:
            if self.path != "/":
                self.send_error(404)
                return
            self._send(200, HTML.encode("utf-8"), "text/html; charset=utf-8")

        def do_POST(self) -> None:
            if self.path != "/api/generate":
                self.send_error(404)
                return
            length = int(self.headers.get("content-length", "0"))
            payload = json.loads(self.rfile.read(length) or b"{}")
            try:
                prompt = str(payload.get("prompt", ""))
                idx = torch.tensor([tokenizer.encode(prompt)], dtype=torch.long, device=device)
                out = model.generate(
                    idx,
                    max_new_tokens=int(payload.get("max_new_tokens", 120)),
                    temperature=float(payload.get("temperature", 0.8)),
                    top_k=int(payload.get("top_k", 20)),
                )
                body = json.dumps({"text": tokenizer.decode(out[0].tolist())}).encode("utf-8")
                self._send(200, body, "application/json")
            except Exception as exc:
                body = json.dumps({"error": str(exc)}).encode("utf-8")
                self._send(400, body, "application/json")

        def log_message(self, format: str, *args: object) -> None:
            return

        def _send(self, status: int, body: bytes, content_type: str) -> None:
            self.send_response(status)
            self.send_header("Content-Type", content_type)
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)

    server = ThreadingHTTPServer((host, port), Handler)
    print(f"Serving playground at http://{host}:{port}")
    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nStopped playground server.")
    finally:
        server.server_close()