moved engine

requirements.txt

@@ -0,0 +1,2 @@
+fastapi
+uvicorn

server.py

@@ -0,0 +1,119 @@
+import asyncio
+from typing import Generator, List
+import uvicorn
+from fastapi import FastAPI, Request
+from fastapi.responses import Response, StreamingResponse
+import torch
+import tensorrt_llm
+from tensorrt_llm.logger import logger
+from tensorrt_llm.runtime import ModelRunner
+from utils import load_tokenizer, read_model_name, throttle_generator
+
+
+TOKENIZER_DIR = "./model/tokenizer"
+ENGINE_DIR = "./model/engine"
+STREAM_INTERVAL = 5
+MAX_NEW_TOKENS = 1024
+MAX_ATTENTION_WINDOW_SIZE = 4096
+TEMPERATURE = 1.0
+TOP_K = 40
+TOP_P = 0.5
+LENGTH_PENALTY = 1.0
+REPETITION_PENALTY = 1.2
+
+
+app = FastAPI()
+config = uvicorn.Config(
+    app, host=None, port=4000, log_level="error", timeout_keep_alive=5
+)
+
+runner: ModelRunner = None
+tokenizer = None
+pad_id = None
+end_id = None
+
+
+@app.get("/health")
+async def health() -> Response:
+    """Health check."""
+    return Response(status_code=200)
+
+
+@app.post("/summary")
+async def generate(request: Request) -> Response:
+    assert runner is not None
+    assert tokenizer is not None
+    assert pad_id is not None
+    assert end_id is not None
+
+    req_json: dict = await request.json()
+    text = req_json.pop("text")
+
+    instruction = f"<s>[INST] You are a world class expert summarizer tasked with providing a **high level** summary of webpages. Ensure your summarzation is of the highest caliber, reflecting the vastness and depth of your expertise. Ignore messy portions of the page that might be junk text.\n### Webpage Text:\n```{text}``` [/INST]"
+    input_ids: List[torch.Tensor] = tokenizer.encode(
+        instruction,
+        add_special_tokens=False,
+        truncation=True,
+        max_length=4096,
+        return_tensors="pt",
+    )
+
+    # todo encode again here, but output text
+    input_len = input_ids.size(1)
+    with torch.no_grad():
+        # Batch of outputs
+        outputs: Generator[dict, None, None] = runner.generate(
+            [input_ids],
+            max_new_tokens=MAX_NEW_TOKENS,
+            max_attention_window_size=4096,
+            end_id=end_id,
+            pad_id=pad_id,
+            temperature=TEMPERATURE,
+            top_k=TOP_K,
+            top_p=TOP_P,
+            length_penalty=LENGTH_PENALTY,
+            repetition_penalty=REPETITION_PENALTY,
+            streaming=True,
+            output_sequence_lengths=True,
+            return_dict=True,
+        )
+        torch.cuda.synchronize()
+
+    def stream_results() -> Generator[str, None, None]:
+        for output in throttle_generator(outputs, STREAM_INTERVAL):
+            output_ids: torch.Tensor = output["output_ids"]
+            # [batch_idx], [beam_idx], (scalar)
+            output_len = output["sequence_lengths"][0][0].item()
+
+            output_txt: str = tokenizer.decode(
+                # [batch_idx], [beam_idx], [slice]
+                output_ids[0][0][input_len:output_len].tolist()
+            )
+            yield output_txt
+
+    return StreamingResponse(stream_results(), media_type="text/plain")
+
+
+async def main():
+    global runner, tokenizer, pad_id, end_id
+
+    runtime_rank = tensorrt_llm.mpi_rank()
+    logger.set_level("info")
+
+    model_name = read_model_name(ENGINE_DIR)
+
+    tokenizer, pad_id, end_id = load_tokenizer(
+        tokenizer_dir=TOKENIZER_DIR,
+        model_name=model_name,
+    )
+
+    runner_kwargs = dict(engine_dir=ENGINE_DIR, rank=runtime_rank, debug_mode=True)
+
+    runner = ModelRunner.from_dir(**runner_kwargs)
+
+    await uvicorn.Server(config).serve()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
+