{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "50532a6d-a3d7-485e-a2cc-76d70cadb0d2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "model_id = \"EleutherAI/gpt-j-6B\"\n",
    "revision = \"float16\"  # use float16 weights to fit in 16GB GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "140cbbc0-9068-4605-ac29-4fe68f4e388a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import ray"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa62af90-6652-4971-8ec1-8f1cac3fb01b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ray.init(\n",
    "    address=\"ray://example-cluster-kuberay-head-svc:10001\",\n",
    "    runtime_env={\n",
    "        \"pip\": [\n",
    "            \"IPython\",\n",
    "            \"boto3==1.26\",\n",
    "            \"botocore==1.29\", \n",
    "            \"datasets\",\n",
    "            \"fastapi\",\n",
    "            \"accelerate>=0.16.0\",\n",
    "            \"transformers>=4.26.0\",\n",
    "            \"numpy<1.24\",  # remove when mlflow updates beyond 2.2\n",
    "            \"torch\",\n",
    "        ]\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30bce290-51ec-4e8e-b6df-c179e8ebc281",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "from ray import serve\n",
    "from starlette.requests import Request\n",
    "\n",
    "\n",
    "@serve.deployment(ray_actor_options={\"num_gpus\": 1})\n",
    "class PredictDeployment:\n",
    "    def __init__(self, model_id: str, revision: str = None):\n",
    "        from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "        import torch\n",
    "\n",
    "        self.model = AutoModelForCausalLM.from_pretrained(\n",
    "            model_id,\n",
    "            revision=revision,\n",
    "            torch_dtype=torch.float16,\n",
    "            low_cpu_mem_usage=True,\n",
    "            device_map=\"auto\",  # automatically makes use of all GPUs available to the Actor\n",
    "        )\n",
    "        self.tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
    "\n",
    "    def generate(self, text: str) -> pd.DataFrame:\n",
    "        input_ids = self.tokenizer(text, return_tensors=\"pt\").input_ids.to(\n",
    "            self.model.device\n",
    "        )\n",
    "\n",
    "        gen_tokens = self.model.generate(\n",
    "            input_ids,\n",
    "            do_sample=True,\n",
    "            temperature=0.9,\n",
    "            max_length=100,\n",
    "        )\n",
    "        return pd.DataFrame(\n",
    "            self.tokenizer.batch_decode(gen_tokens), columns=[\"responses\"]\n",
    "        )\n",
    "\n",
    "    async def __call__(self, http_request: Request) -> str:\n",
    "        json_request: str = await http_request.json()\n",
    "        prompts = []\n",
    "        for prompt in json_request:\n",
    "            text = prompt[\"text\"]\n",
    "            if isinstance(text, list):\n",
    "                prompts.extend(text)\n",
    "            else:\n",
    "                prompts.append(text)\n",
    "        return self.generate(prompts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33ea1b4a-afcf-4717-8ba7-3375331381ba",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "deployment = PredictDeployment.bind(model_id=model_id, revision=revision)\n",
    "serve.run(deployment, host=\"0.0.0.0\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bce2ad16-5248-4760-b92c-cb4517cb110c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "prompt = (\n",
    "    \"In a shocking finding, scientists discovered a herd of unicorns living in a remote, \"\n",
    "    \"previously unexplored valley, in the Andes Mountains. Even more surprising to the \"\n",
    "    \"researchers was the fact that the unicorns spoke perfect English.\"\n",
    ")\n",
    "\n",
    "sample_input = {\"text\": prompt}\n",
    "\n",
    "output = requests.post(\"http://example-cluster-kuberay-head-svc:8000/\", json=[sample_input]).json()\n",
    "print(output)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96f74370-2ac1-4dcd-81a5-e95cd2b8ac51",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}