Downloads: 1
309 lines (308 with data), 11.4 kB
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "8ae13342-5979-42df-9fea-5bb97c60be23",
"metadata": {},
"outputs": [],
"source": [
"import sys\n",
"import os\n",
"import time\n",
"from pathlib import Path\n",
"from typing import Any, Literal, Optional\n",
"import json\n",
"\n",
"import lightning as L\n",
"import torch\n",
"import torch._dynamo.config\n",
"import torch._inductor.config\n",
"from lightning.fabric.plugins import BitsandbytesPrecision\n",
"from lightning.fabric.strategies import FSDPStrategy\n",
"\n",
"## Add the lit_gpt folder to the path\n",
"sys.path.insert(0, os.path.abspath('../'))\n",
"\n",
"from lit_gpt import GPT, Config, Tokenizer\n",
"from lit_gpt.model import Block\n",
"from lit_gpt.utils import (\n",
" check_valid_checkpoint_dir,\n",
" get_default_supported_precision,\n",
" gptq_quantization,\n",
" load_checkpoint,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7ac87e9a-6846-4f84-9c02-27f46fa417ce",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"torch.set_float32_matmul_precision(\"high\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e6eec0b2-81d9-4a88-b20c-0be57ec6b2a2",
"metadata": {},
"outputs": [],
"source": [
"with open('../data/entity_extraction/entity-extraction-test-data.json', 'r') as file:\n",
" test_data = json.load(file)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4d92ad07",
"metadata": {},
"outputs": [],
"source": [
"example = {\n",
" \"input\": \"Natalie Cooper,\\nncooper@example.com\\n6789 Birch Street, Denver, CO 80203,\\n303-555-6543, United States\\n\\nRelationship to XYZ Pharma Inc.: Patient\\nReason for contacting: Adverse Event\\n\\nMessage: Hi, after starting Abilify for bipolar I disorder, I've noticed that I am experiencing nausea and vomiting. Are these typical reactions? Best, Natalie Cooper\",\n",
" \"output\": \"{\\\"drug_name\\\": \\\"Abilify\\\", \\\"adverse_events\\\": [\\\"nausea\\\", \\\"vomiting\\\"]}\"\n",
" }"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3d5eed73-7020-466a-8b83-5ebc7d497698",
"metadata": {},
"outputs": [],
"source": [
"prompt = f\"\"\"Act as an expert Analyst with 20+ years of experience\\\n",
"in Pharma and Healthcare industry. \\\n",
"For the following provided input you need to generate the output which \\\n",
"identifies and extracts entities like 'drug_name' and 'adverse_events' \\\n",
"use the format:\\n\\\n",
"{{'drug_name':'DRUG_NAME_HERE', 'adverse_events':[## List of symptoms here]}}\\n\\\n",
"\n",
"### Extract Entities from the follwing:\\n\\\n",
"{example[\"input\"]}\\\n",
"\n",
"### Response:\n",
"\"\"\""
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c972c1d6-f2f3-4bf8-bd92-405bda4b02ab",
"metadata": {},
"outputs": [],
"source": [
"print(prompt)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3ae2d36b-3f58-4be1-955d-cca8c8c5f7d6",
"metadata": {},
"outputs": [],
"source": [
"def multinomial_num_samples_1(probs: torch.Tensor) -> torch.Tensor:\n",
" if torch._dynamo.is_compiling():\n",
" # Faster alternative to `torch.multinomial(probs, num_samples=1)` that is also CUDAGraph friendly\n",
" distribution = torch.empty_like(probs).exponential_(1)\n",
" return torch.argmax(probs / distribution, dim=-1, keepdim=True)\n",
" return torch.multinomial(probs, num_samples=1)\n",
"\n",
"\n",
"def sample(logits: torch.Tensor, temperature: float = 1.0, top_k: Optional[int] = None) -> torch.Tensor:\n",
" logits = logits[0, -1]\n",
" # optionally crop the logits to only the top k options\n",
" if top_k is not None:\n",
" v, i = torch.topk(logits, min(top_k, logits.size(-1)))\n",
" # do not use `torch.where` as in nanogpt because it will repeat top-k collisions\n",
" logits = torch.full_like(logits, float(\"-inf\")).scatter_(-1, i, v)\n",
" # optionally scale the logits and sample from a probability distribution\n",
" if temperature > 0.0:\n",
" probs = torch.nn.functional.softmax(logits / temperature, dim=-1)\n",
" return multinomial_num_samples_1(probs)\n",
" return torch.argmax(logits, dim=-1, keepdim=True)\n",
"\n",
"def next_token(model: GPT, input_pos: torch.Tensor, x: torch.Tensor, **kwargs: Any) -> torch.Tensor:\n",
" logits = model(x, input_pos)\n",
" next = sample(logits, **kwargs)\n",
" return next.type_as(x)\n",
"\n",
"@torch.inference_mode()\n",
"def generate(\n",
" model: GPT,\n",
" prompt: torch.Tensor,\n",
" max_returned_tokens: int,\n",
" *,\n",
" temperature: float = 1.0,\n",
" top_k: Optional[int] = None,\n",
" eos_id: Optional[int] = None,\n",
") -> torch.Tensor:\n",
" \"\"\"Takes a conditioning sequence (prompt) as input and continues to generate as many tokens as requested.\n",
"\n",
" The implementation of this function is modified from A. Karpathy's nanoGPT.\n",
"\n",
" Args:\n",
" model: The model to use.\n",
" prompt: Tensor of shape (T) with indices of the prompt sequence.\n",
" max_returned_tokens: The maximum number of tokens to return (given plus generated).\n",
" temperature: Scales the predicted logits by 1 / temperature.\n",
" top_k: If specified, only sample among the tokens with the k highest probabilities.\n",
" eos_id: If specified, stop generating any more token once the <eos> token is triggered.\n",
" \"\"\"\n",
" T = prompt.size(0)\n",
" assert max_returned_tokens > T\n",
" if model.max_seq_length < max_returned_tokens - 1:\n",
" # rolling the kv cache based on the `input_pos` value would be necessary. However, doing so would introduce a\n",
" # data dependency on the `input_pos` tensor and impact model compilation. Since this setting is uncommon, we do\n",
" # not support it to avoid negatively impacting the overall speed\n",
" raise NotImplementedError(f\"max_seq_length {model.max_seq_length} needs to be >= {max_returned_tokens - 1}\")\n",
"\n",
" device = prompt.device\n",
" tokens = [prompt]\n",
" input_pos = torch.tensor([T], device=device)\n",
" token = next_token(\n",
" model, torch.arange(0, T, device=device), prompt.view(1, -1), temperature=temperature, top_k=top_k\n",
" ).clone()\n",
" tokens.append(token)\n",
" for _ in range(2, max_returned_tokens - T + 1):\n",
" token = next_token(model, input_pos, token.view(1, -1), temperature=temperature, top_k=top_k).clone()\n",
" tokens.append(token)\n",
" if token == eos_id:\n",
" break\n",
" input_pos = input_pos.add_(1)\n",
" return torch.cat(tokens)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e3c30b09-e631-44b8-a45b-ac79dbbab52f",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"print(\"[INFO] Using StableLM-3B base model\")\n",
"checkpoint_dir: Path = Path(\"../checkpoints/stabilityai/stablelm-base-alpha-3b\")\n",
"predictions_file_name = '../data/predictions-stablelm-base.json'\n",
"\n",
"quantize: Optional[Literal[\"bnb.nf4\", \"bnb.nf4-dq\", \"bnb.fp4\", \"bnb.fp4-dq\", \"bnb.int8\", \"gptq.int4\"]] = None\n",
"max_new_tokens: int = 50\n",
"top_k: int = 200\n",
"temperature: float = 0.1\n",
"strategy: str = \"auto\"\n",
"devices: int = 1\n",
"precision: Optional[str] = None\n",
"num_samples: int = 1,"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "b842b257-2c9c-4c64-b011-f97e215c3794",
"metadata": {},
"outputs": [],
"source": [
"if strategy == \"fsdp\":\n",
" strategy = FSDPStrategy(auto_wrap_policy={Block}, cpu_offload=False)\n",
"fabric = L.Fabric(devices=devices, precision=precision, strategy=strategy)\n",
"fabric.launch()\n",
"\n",
"check_valid_checkpoint_dir(checkpoint_dir)\n",
"\n",
"config = Config.from_json(checkpoint_dir / \"lit_config.json\")\n",
"\n",
"if quantize is not None and devices > 1:\n",
" raise NotImplementedError\n",
"if quantize == \"gptq.int4\":\n",
" model_file = \"lit_model_gptq.4bit.pth\"\n",
" if not (checkpoint_dir / model_file).is_file():\n",
" raise ValueError(\"Please run `python quantize/gptq.py` first\")\n",
"else:\n",
" model_file = \"lit_model.pth\"\n",
" \n",
"checkpoint_path = checkpoint_dir / model_file\n",
"tokenizer = Tokenizer(checkpoint_dir)\n",
"encoded = tokenizer.encode(prompt, device=fabric.device)\n",
"prompt_length = encoded.size(0)\n",
"max_returned_tokens = prompt_length + max_new_tokens\n",
"\n",
"fabric.print(f\"Loading model {str(checkpoint_path)!r} with {config.__dict__}\", file=sys.stderr)\n",
"t0 = time.perf_counter()\n",
"with fabric.init_module(empty_init=True), gptq_quantization(quantize == \"gptq.int4\"):\n",
" model = GPT(config)\n",
"fabric.print(f\"Time to instantiate model: {time.perf_counter() - t0:.02f} seconds.\", file=sys.stderr)\n",
"with fabric.init_tensor():\n",
" # set the max_seq_length to limit the memory usage to what we need\n",
" model.max_seq_length = max_returned_tokens\n",
" # enable the kv cache\n",
" model.set_kv_cache(batch_size=1)\n",
"model.eval()\n",
"\n",
"model = fabric.setup_module(model)\n",
"t0 = time.perf_counter()\n",
"load_checkpoint(fabric, model, checkpoint_path)\n",
"fabric.print(f\"Time to load the model weights: {time.perf_counter() - t0:.02f} seconds.\", file=sys.stderr)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "59e00f69-d9fd-4cf5-a4ed-f142991b870f",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"prompt"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3eeaa237-7a95-4c9d-aeda-8a3c246d02fe",
"metadata": {},
"outputs": [],
"source": [
"L.seed_everything(1234)\n",
"\n",
"t0 = time.perf_counter()\n",
"y = generate(model, encoded, max_returned_tokens, temperature=temperature, top_k=top_k)\n",
"t = time.perf_counter() - t0\n",
"for block in model.transformer.h:\n",
" block.attn.kv_cache.reset_parameters()\n",
"output = tokenizer.decode(y)\n",
"fabric.print(output)\n",
"tokens_generated = y.size(0) - prompt_length"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "new",
"language": "python",
"name": "new"
},
"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.13"
}
},
"nbformat": 4,
"nbformat_minor": 5
}