mirror of
https://github.com/wassname/prob_jsonformer.git
synced 2026-08-20 12:40:20 +08:00
working
This commit is contained in:
wassname
@@ -1,246 +0,0 @@
|
||||
{
|
||||
"cells": [
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {
|
||||
"id": "view-in-github",
|
||||
"colab_type": "text"
|
||||
},
|
||||
"source": [
|
||||
"<a href=\"https://colab.research.google.com/github/1rgs/jsonformer/blob/main/Jsonformer_example.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "p-kogdhHO3PC",
|
||||
"outputId": "d41737a8-998e-4451-cce3-6ac6c98cc69b"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Looking in indexes: https://pypi.org/simple, https://us-python.pkg.dev/colab-wheels/public/simple/\n",
|
||||
"Requirement already satisfied: transformers in /usr/local/lib/python3.10/dist-packages (4.28.1)\n",
|
||||
"Requirement already satisfied: accelerate in /usr/local/lib/python3.10/dist-packages (0.18.0)\n",
|
||||
"Requirement already satisfied: jsonformer in /usr/local/lib/python3.10/dist-packages (0.9.0)\n",
|
||||
"Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.10/dist-packages (from transformers) (4.65.0)\n",
|
||||
"Requirement already satisfied: huggingface-hub<1.0,>=0.11.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.14.1)\n",
|
||||
"Requirement already satisfied: requests in /usr/local/lib/python3.10/dist-packages (from transformers) (2.27.1)\n",
|
||||
"Requirement already satisfied: packaging>=20.0 in /usr/local/lib/python3.10/dist-packages (from transformers) (23.1)\n",
|
||||
"Requirement already satisfied: pyyaml>=5.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (6.0)\n",
|
||||
"Requirement already satisfied: numpy>=1.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (1.22.4)\n",
|
||||
"Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.10/dist-packages (from transformers) (2022.10.31)\n",
|
||||
"Requirement already satisfied: filelock in /usr/local/lib/python3.10/dist-packages (from transformers) (3.12.0)\n",
|
||||
"Requirement already satisfied: tokenizers!=0.11.3,<0.14,>=0.11.1 in /usr/local/lib/python3.10/dist-packages (from transformers) (0.13.3)\n",
|
||||
"Requirement already satisfied: torch>=1.4.0 in /usr/local/lib/python3.10/dist-packages (from accelerate) (2.0.0+cu118)\n",
|
||||
"Requirement already satisfied: psutil in /usr/local/lib/python3.10/dist-packages (from accelerate) (5.9.5)\n",
|
||||
"Requirement already satisfied: termcolor<3.0.0,>=2.3.0 in /usr/local/lib/python3.10/dist-packages (from jsonformer) (2.3.0)\n",
|
||||
"Requirement already satisfied: fsspec in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.11.0->transformers) (2023.4.0)\n",
|
||||
"Requirement already satisfied: typing-extensions>=3.7.4.3 in /usr/local/lib/python3.10/dist-packages (from huggingface-hub<1.0,>=0.11.0->transformers) (4.5.0)\n",
|
||||
"Requirement already satisfied: sympy in /usr/local/lib/python3.10/dist-packages (from torch>=1.4.0->accelerate) (1.11.1)\n",
|
||||
"Requirement already satisfied: networkx in /usr/local/lib/python3.10/dist-packages (from torch>=1.4.0->accelerate) (3.1)\n",
|
||||
"Requirement already satisfied: triton==2.0.0 in /usr/local/lib/python3.10/dist-packages (from torch>=1.4.0->accelerate) (2.0.0)\n",
|
||||
"Requirement already satisfied: jinja2 in /usr/local/lib/python3.10/dist-packages (from torch>=1.4.0->accelerate) (3.1.2)\n",
|
||||
"Requirement already satisfied: lit in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.4.0->accelerate) (16.0.2)\n",
|
||||
"Requirement already satisfied: cmake in /usr/local/lib/python3.10/dist-packages (from triton==2.0.0->torch>=1.4.0->accelerate) (3.25.2)\n",
|
||||
"Requirement already satisfied: idna<4,>=2.5 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (3.4)\n",
|
||||
"Requirement already satisfied: urllib3<1.27,>=1.21.1 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (1.26.15)\n",
|
||||
"Requirement already satisfied: charset-normalizer~=2.0.0 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2.0.12)\n",
|
||||
"Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.10/dist-packages (from requests->transformers) (2022.12.7)\n",
|
||||
"Requirement already satisfied: MarkupSafe>=2.0 in /usr/local/lib/python3.10/dist-packages (from jinja2->torch>=1.4.0->accelerate) (2.1.2)\n",
|
||||
"Requirement already satisfied: mpmath>=0.19 in /usr/local/lib/python3.10/dist-packages (from sympy->torch>=1.4.0->accelerate) (1.3.0)\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"!pip install transformers accelerate jsonformer"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "9y0Jjj0sOtor",
|
||||
"outputId": "6c5ea465-384f-4aa1-8304-6fa721534f73"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Loading model and tokenizer...\n",
|
||||
"Loaded model and tokenizer\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from transformers import AutoModelForCausalLM, AutoTokenizer\n",
|
||||
"\n",
|
||||
"print(\"Loading model and tokenizer...\")\n",
|
||||
"model_name = \"databricks/dolly-v2-3b\"\n",
|
||||
"model = AutoModelForCausalLM.from_pretrained(model_name, use_cache=True, device_map=\"auto\")\n",
|
||||
"tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, use_cache=True)\n",
|
||||
"print(\"Loaded model and tokenizer\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 6,
|
||||
"metadata": {
|
||||
"colab": {
|
||||
"base_uri": "https://localhost:8080/"
|
||||
},
|
||||
"id": "8E7zFQ3NO0Yd",
|
||||
"outputId": "5f4d1ad0-1acf-48e0-d9bd-82f88ebcf97a"
|
||||
},
|
||||
"outputs": [
|
||||
{
|
||||
"output_type": "stream",
|
||||
"name": "stdout",
|
||||
"text": [
|
||||
"Generating...\n",
|
||||
"{\n",
|
||||
" car: {\n",
|
||||
" make: \"audi\",\n",
|
||||
" model: \"model A8\",\n",
|
||||
" year: 2016.0,\n",
|
||||
" colors: [\n",
|
||||
" \"blue\"\n",
|
||||
" ],\n",
|
||||
" features: {\n",
|
||||
" audio: {\n",
|
||||
" brand: \"sony\",\n",
|
||||
" speakers: 2.0,\n",
|
||||
" hasBluetooth: True\n",
|
||||
" },\n",
|
||||
" safety: {\n",
|
||||
" airbags: 2.0,\n",
|
||||
" parkingSensors: True,\n",
|
||||
" laneAssist: True\n",
|
||||
" },\n",
|
||||
" performance: {\n",
|
||||
" engine: \"4.0\",\n",
|
||||
" horsepower: 220.0,\n",
|
||||
" topSpeed: 220.0\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" },\n",
|
||||
" owner: {\n",
|
||||
" firstName: \"John\",\n",
|
||||
" lastName: \"Doe\",\n",
|
||||
" age: 40.0\n",
|
||||
" }\n",
|
||||
"}\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from jsonformer.format import highlight_values\n",
|
||||
"from jsonformer.main import Jsonformer\n",
|
||||
"\n",
|
||||
"car = {\n",
|
||||
" \"type\": \"object\",\n",
|
||||
" \"properties\": {\n",
|
||||
" \"car\": {\n",
|
||||
" \"type\": \"object\",\n",
|
||||
" \"properties\": {\n",
|
||||
" \"make\": {\"type\": \"string\"},\n",
|
||||
" \"model\": {\"type\": \"string\"},\n",
|
||||
" \"year\": {\"type\": \"number\"},\n",
|
||||
" \"colors\": {\n",
|
||||
" \"type\": \"array\",\n",
|
||||
" \"items\": {\"type\": \"string\"}\n",
|
||||
" },\n",
|
||||
" \"features\": {\n",
|
||||
" \"type\": \"object\",\n",
|
||||
" \"properties\": {\n",
|
||||
" \"audio\": {\n",
|
||||
" \"type\": \"object\",\n",
|
||||
" \"properties\": {\n",
|
||||
" \"brand\": {\"type\": \"string\"},\n",
|
||||
" \"speakers\": {\"type\": \"number\"},\n",
|
||||
" \"hasBluetooth\": {\"type\": \"boolean\"}\n",
|
||||
" }\n",
|
||||
" },\n",
|
||||
" \"safety\": {\n",
|
||||
" \"type\": \"object\",\n",
|
||||
" \"properties\": {\n",
|
||||
" \"airbags\": {\"type\": \"number\"},\n",
|
||||
" \"parkingSensors\": {\"type\": \"boolean\"},\n",
|
||||
" \"laneAssist\": {\"type\": \"boolean\"}\n",
|
||||
" }\n",
|
||||
" },\n",
|
||||
" \"performance\": {\n",
|
||||
" \"type\": \"object\",\n",
|
||||
" \"properties\": {\n",
|
||||
" \"engine\": {\"type\": \"string\"},\n",
|
||||
" \"horsepower\": {\"type\": \"number\"},\n",
|
||||
" \"topSpeed\": {\"type\": \"number\"}\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" },\n",
|
||||
" \"owner\": {\n",
|
||||
" \"type\": \"object\",\n",
|
||||
" \"properties\": {\n",
|
||||
" \"firstName\": {\"type\": \"string\"},\n",
|
||||
" \"lastName\": {\"type\": \"string\"},\n",
|
||||
" \"age\": {\"type\": \"number\"},\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
" }\n",
|
||||
"}\n",
|
||||
"\n",
|
||||
"builder = Jsonformer(\n",
|
||||
" model=model,\n",
|
||||
" tokenizer=tokenizer,\n",
|
||||
" json_schema=car,\n",
|
||||
" prompt=\"Generate an example car\",\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"print(\"Generating...\")\n",
|
||||
"output = builder()\n",
|
||||
"\n",
|
||||
"highlight_values(output)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"source": [],
|
||||
"metadata": {
|
||||
"id": "0eA8hzFHA1jo"
|
||||
},
|
||||
"execution_count": null,
|
||||
"outputs": []
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
"accelerator": "GPU",
|
||||
"colab": {
|
||||
"machine_shape": "hm",
|
||||
"provenance": [],
|
||||
"include_colab_link": true
|
||||
},
|
||||
"gpuClass": "premium",
|
||||
"kernelspec": {
|
||||
"display_name": "Python 3",
|
||||
"name": "python3"
|
||||
},
|
||||
"language_info": {
|
||||
"name": "python"
|
||||
}
|
||||
},
|
||||
"nbformat": 4,
|
||||
"nbformat_minor": 0
|
||||
}
|
||||
+118
-159
@@ -4,6 +4,17 @@
|
||||
"cell_type": "code",
|
||||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# autoreload your package\n",
|
||||
"%load_ext autoreload\n",
|
||||
"%autoreload 2"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 2,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
@@ -28,7 +39,7 @@
|
||||
" warnings.warn(\n",
|
||||
"`flash-attention` package not found, consider installing for better performance: No module named 'flash_attn'.\n",
|
||||
"Current `flash-attenton` does not support `window_size`. Either upgrade or use `attn_implementation='eager'`.\n",
|
||||
"Loading checkpoint shards: 100%|██████████| 4/4 [00:04<00:00, 1.10s/it]\n",
|
||||
"Loading checkpoint shards: 100%|██████████| 4/4 [00:02<00:00, 1.94it/s]\n",
|
||||
"/media/wassname/SGIronWolf/projects5/2024/prob_jsonformer/.venv/lib/python3.9/site-packages/huggingface_hub/file_download.py:1132: FutureWarning: `resume_download` is deprecated and will be removed in version 1.0.0. Downloads always resume when possible. If you want to force a new download, use `force_download=True`.\n",
|
||||
" warnings.warn(\n",
|
||||
"Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
|
||||
@@ -44,6 +55,7 @@
|
||||
],
|
||||
"source": [
|
||||
"from transformers import AutoModelForCausalLM, AutoTokenizer\n",
|
||||
"import torch\n",
|
||||
"\n",
|
||||
"print(\"Loading model and tokenizer...\")\n",
|
||||
"# model_name = \"databricks/dolly-v2-3b\"\n",
|
||||
@@ -51,8 +63,10 @@
|
||||
"model = AutoModelForCausalLM.from_pretrained(\n",
|
||||
" model_name,\n",
|
||||
" use_cache=True,\n",
|
||||
" torch_dtype=torch.float16,\n",
|
||||
" # device=\"cuda:0\",\n",
|
||||
" # device_map=\"auto\",\n",
|
||||
" attn_implementation='eager',\n",
|
||||
" trust_remote_code=True,\n",
|
||||
").to(\"cuda:0\")\n",
|
||||
"tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True, use_cache=True)\n",
|
||||
@@ -68,189 +82,118 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 65,
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from jaxtyping import Float, Int\n",
|
||||
"import torch\n",
|
||||
"from torch.nn import functional as F\n",
|
||||
"from torch import Tensor\n",
|
||||
"from typing import List, Callable, Tuple, Dict, Optional\n",
|
||||
"import pandas as pd"
|
||||
"# from jaxtyping import Float, Int\n",
|
||||
"# import torch\n",
|
||||
"# from torch.nn import functional as F\n",
|
||||
"# from torch import Tensor\n",
|
||||
"# from typing import List, Callable, Tuple, Dict, Optional\n",
|
||||
"# import pandas as pd"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 43,
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"[tensor([ 1, 29871, 29896]),\n",
|
||||
" tensor([ 1, 278, 18109, 11285]),\n",
|
||||
" tensor([ 1, 278, 289, 4992, 12544])]"
|
||||
]
|
||||
},
|
||||
"execution_count": 43,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# initital state\n",
|
||||
"prompt = \"The necromancer in his tower, what's his top problem? \"\n",
|
||||
"choices = [\"1\", \"the skeleton\", \"the boney boys\"]\n",
|
||||
"choices_tokens = tokenizer(choices).input_ids\n",
|
||||
"choices_tokens = [torch.tensor(c) for c in choices_tokens]\n",
|
||||
"# current_tokens = torch.tensor([])\n",
|
||||
"# # initital state\n",
|
||||
"# prompt = \"The necromancer in his tower, what's his top problem? \"\n",
|
||||
"# choices = [\"1\", \"the skeleton\", \"the boney boys\"]\n",
|
||||
"# choices_tokens = tokenizer(choices).input_ids\n",
|
||||
"# choices_tokens = [torch.tensor(c) for c in choices_tokens]\n",
|
||||
"# # current_tokens = torch.tensor([])\n",
|
||||
"\n",
|
||||
"# next\n",
|
||||
"input_ids = tokenizer([prompt], return_tensors=\"pt\").to(model.device).input_ids[0]\n",
|
||||
"choices_tokens\n",
|
||||
"# # next\n",
|
||||
"# input_ids = tokenizer([prompt], return_tensors=\"pt\").to(model.device).input_ids[0]\n",
|
||||
"# choices_tokens\n",
|
||||
"\n",
|
||||
"# for each next choice, continue down the tree, recording the log probs"
|
||||
"# # for each next choice, continue down the tree, recording the log probs"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 70,
|
||||
"execution_count": 5,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"tensor([1]) current_tokens\n",
|
||||
"tensor([ 1, 278]) current_tokens\n",
|
||||
"tensor([ 1, 278, 289]) current_tokens\n",
|
||||
"tensor([ 1, 278, 289, 4992]) current_tokens\n",
|
||||
"tensor([ 1, 278, 289, 4992, 12544]) current_tokens\n",
|
||||
"tensor([ 1, 278, 18109]) current_tokens\n",
|
||||
"tensor([ 1, 278, 18109, 11285]) current_tokens\n",
|
||||
"tensor([ 1, 29871]) current_tokens\n",
|
||||
"tensor([ 1, 29871, 29896]) current_tokens\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"def get_valid_next_choices(choices_tokens, current_tokens):\n",
|
||||
" next_choices = []\n",
|
||||
" for choice_tokens in choices_tokens:\n",
|
||||
" # if we have some more slots left\n",
|
||||
" if len(current_tokens)<len(choice_tokens):\n",
|
||||
" # see if current_tokens matches\n",
|
||||
" if (choice_tokens[: len(current_tokens)] == current_tokens).all():\n",
|
||||
" c = choice_tokens[len(current_tokens)].item()\n",
|
||||
" next_choices.append(c)\n",
|
||||
"# def get_valid_next_choices(choices_tokens, current_tokens):\n",
|
||||
"# next_choices = []\n",
|
||||
"# for choice_tokens in choices_tokens:\n",
|
||||
"# # if we have some more slots left\n",
|
||||
"# if len(current_tokens)<len(choice_tokens):\n",
|
||||
"# # see if current_tokens matches\n",
|
||||
"# if (choice_tokens[: len(current_tokens)] == current_tokens).all():\n",
|
||||
"# c = choice_tokens[len(current_tokens)].item()\n",
|
||||
"# next_choices.append(c)\n",
|
||||
"\n",
|
||||
" next_choices = list(set(next_choices))\n",
|
||||
" return torch.LongTensor(next_choices)\n",
|
||||
"# next_choices = list(set(next_choices))\n",
|
||||
"# return torch.LongTensor(next_choices)\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"def next(\n",
|
||||
" input_ids: Int[Tensor, \"seq\"],\n",
|
||||
" choice: Optional[Int[Tensor, \"\"]] = None,\n",
|
||||
" prob: float = 1,\n",
|
||||
" current_tokens: Int[Tensor, \"seq\"] = torch.LongTensor([]),\n",
|
||||
" z=[],\n",
|
||||
"):\n",
|
||||
" if choice is not None:\n",
|
||||
" c = choice[None].to(current_tokens.device)\n",
|
||||
" current_tokens = torch.cat([current_tokens, c], dim=-1)\n",
|
||||
" print(current_tokens, 'current_tokens')\n",
|
||||
" c = choice[None].to(input_ids.device)\n",
|
||||
" input_ids = torch.cat([input_ids, c], dim=-1)\n",
|
||||
"# def gen_choice_probs(\n",
|
||||
"# model: AutoModelForCausalLM,\n",
|
||||
"# tokenizer: AutoTokenizer,\n",
|
||||
"# input_ids: Int[Tensor, \"seq\"],\n",
|
||||
"# choices_tokens: List[Int[Tensor, \"seq\"]],\n",
|
||||
"# choice: Optional[Int[Tensor, \"\"]] = None,\n",
|
||||
"# prob: float = 1,\n",
|
||||
"# current_tokens: Int[Tensor, \"seq\"] = torch.LongTensor([]),\n",
|
||||
"# z=[],\n",
|
||||
"# ):\n",
|
||||
"# if choice is not None:\n",
|
||||
"# c = choice[None].to(current_tokens.device)\n",
|
||||
"# current_tokens = torch.cat([current_tokens, c], dim=-1)\n",
|
||||
"# print(current_tokens, 'current_tokens')\n",
|
||||
"# c = choice[None].to(input_ids.device)\n",
|
||||
"# input_ids = torch.cat([input_ids, c], dim=-1)\n",
|
||||
"\n",
|
||||
" next_choices = get_valid_next_choices(choices_tokens, current_tokens)\n",
|
||||
" if len(next_choices) == 0:\n",
|
||||
" s = tokenizer.decode(current_tokens)\n",
|
||||
" r = dict(tokens=current_tokens.cpu(), prob=prob, choice=s)\n",
|
||||
" yield r\n",
|
||||
" else:\n",
|
||||
" o = model(input_ids[None])\n",
|
||||
" logits_constrained = o.logits[0, -1][next_choices]\n",
|
||||
" probs = F.softmax(logits_constrained, dim=-1)\n",
|
||||
" for i in range(len(next_choices)):\n",
|
||||
" next_choice = next_choices[i]\n",
|
||||
" next_prob = prob * probs[i].item()\n",
|
||||
" yield from next(\n",
|
||||
" input_ids=input_ids,\n",
|
||||
" choice=next_choice,\n",
|
||||
" prob=next_prob,\n",
|
||||
" current_tokens=current_tokens,\n",
|
||||
" z=z + [i],\n",
|
||||
" )\n",
|
||||
"# next_choices = get_valid_next_choices(choices_tokens, current_tokens)\n",
|
||||
"# if len(next_choices) == 0:\n",
|
||||
"# s = tokenizer.decode(current_tokens)\n",
|
||||
"# r = dict(tokens=current_tokens.cpu(), prob=prob, choice=s)\n",
|
||||
"# yield r\n",
|
||||
"# else:\n",
|
||||
"# o = model(input_ids[None])\n",
|
||||
"# logits_constrained = o.logits[0, -1][next_choices]\n",
|
||||
"# probs = F.softmax(logits_constrained, dim=-1)\n",
|
||||
"# for i in range(len(next_choices)):\n",
|
||||
"# next_choice = next_choices[i]\n",
|
||||
"# next_prob = prob * probs[i].item()\n",
|
||||
"# yield from gen_choice_probs(\n",
|
||||
"# model=model,\n",
|
||||
"# tokenizer=tokenizer,\n",
|
||||
"# choices_tokens=choices_tokens,\n",
|
||||
"# input_ids=input_ids,\n",
|
||||
"# choice=next_choice,\n",
|
||||
"# prob=next_prob,\n",
|
||||
"# current_tokens=current_tokens,\n",
|
||||
"# z=z + [i],\n",
|
||||
"# )\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"r = list(next(input_ids=input_ids))"
|
||||
"# r = list(gen_choice_probs(model, tokenizer, input_ids, choices_tokens))"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 74,
|
||||
"execution_count": 6,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"data": {
|
||||
"text/html": [
|
||||
"<div>\n",
|
||||
"<style scoped>\n",
|
||||
" .dataframe tbody tr th:only-of-type {\n",
|
||||
" vertical-align: middle;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" .dataframe tbody tr th {\n",
|
||||
" vertical-align: top;\n",
|
||||
" }\n",
|
||||
"\n",
|
||||
" .dataframe thead th {\n",
|
||||
" text-align: right;\n",
|
||||
" }\n",
|
||||
"</style>\n",
|
||||
"<table border=\"1\" class=\"dataframe\">\n",
|
||||
" <thead>\n",
|
||||
" <tr style=\"text-align: right;\">\n",
|
||||
" <th></th>\n",
|
||||
" <th>prob</th>\n",
|
||||
" <th>choice</th>\n",
|
||||
" </tr>\n",
|
||||
" </thead>\n",
|
||||
" <tbody>\n",
|
||||
" <tr>\n",
|
||||
" <th>2</th>\n",
|
||||
" <td>0.995732</td>\n",
|
||||
" <td><s> 1</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>0</th>\n",
|
||||
" <td>0.004187</td>\n",
|
||||
" <td><s> the boney boys</td>\n",
|
||||
" </tr>\n",
|
||||
" <tr>\n",
|
||||
" <th>1</th>\n",
|
||||
" <td>0.000081</td>\n",
|
||||
" <td><s> the skeleton</td>\n",
|
||||
" </tr>\n",
|
||||
" </tbody>\n",
|
||||
"</table>\n",
|
||||
"</div>"
|
||||
],
|
||||
"text/plain": [
|
||||
" prob choice\n",
|
||||
"2 0.995732 <s> 1\n",
|
||||
"0 0.004187 <s> the boney boys\n",
|
||||
"1 0.000081 <s> the skeleton"
|
||||
]
|
||||
},
|
||||
"execution_count": 74,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"pd.DataFrame(r).sort_values(\"prob\", ascending=False).drop(columns=[\"tokens\"])"
|
||||
"# r"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"# pd.DataFrame(r).sort_values(\"prob\", ascending=False).drop(columns=[\"tokens\"])"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -262,9 +205,24 @@
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": null,
|
||||
"execution_count": 8,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Generating...\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"You are not running the flash-attention implementation, expect numerical differences.\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"from prob_jsonformer.format import highlight_values\n",
|
||||
"from prob_jsonformer.main import Jsonformer\n",
|
||||
@@ -277,6 +235,7 @@
|
||||
" \"properties\": {\n",
|
||||
" \"name\": {\"type\": \"string\"},\n",
|
||||
" \"location\": {\"type\": \"string\"},\n",
|
||||
" \"choices\": {\"type\": \"choices\", \"enum\": [\"1\", \"the\", \"they walked the old dog\", \"1 the\"]},\n",
|
||||
" \"inventory\": {\n",
|
||||
" \"type\": \"array\",\n",
|
||||
" \"items\": {\n",
|
||||
@@ -286,9 +245,9 @@
|
||||
" \"name\": {\"type\": \"string\"},\n",
|
||||
" \"description\": {\"type\": \"string\"},\n",
|
||||
" \"category\": {\"type\": \"string\"},\n",
|
||||
" \"price\": {\"type\": \"number\"},\n",
|
||||
" # \"price\": {\"type\": \"number\"},\n",
|
||||
" \"inStock\": {\"type\": \"boolean\"},\n",
|
||||
" \"rating\": {\"type\": \"number\"},\n",
|
||||
" # \"rating\": {\"type\": \"number\"},\n",
|
||||
" \"images\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}},\n",
|
||||
" },\n",
|
||||
" },\n",
|
||||
|
||||
Generated
+1272
-1012
File diff suppressed because it is too large
Load Diff
@@ -1,2 +1,2 @@
|
||||
from prob_jsonformer.main import prob_jsonformer
|
||||
from prob_jsonformer.main import Jsonformer
|
||||
from prob_jsonformer.format import highlight_values
|
||||
@@ -0,0 +1,61 @@
|
||||
from jaxtyping import Float, Int
|
||||
import torch
|
||||
from torch.nn import functional as F
|
||||
from torch import Tensor
|
||||
from typing import List, Callable, Tuple, Dict, Optional
|
||||
import pandas as pd
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
|
||||
def get_valid_next_choices(choices_tokens, current_tokens):
|
||||
next_choices = []
|
||||
for choice_tokens in choices_tokens:
|
||||
# if we have some more slots left
|
||||
if len(current_tokens) < len(choice_tokens):
|
||||
# see if current_tokens matches
|
||||
if (choice_tokens[: len(current_tokens)] == current_tokens).all():
|
||||
c = choice_tokens[len(current_tokens)].item()
|
||||
next_choices.append(c)
|
||||
|
||||
next_choices = list(set(next_choices))
|
||||
return torch.LongTensor(next_choices)
|
||||
|
||||
|
||||
def choice_tree(
|
||||
model: AutoModelForCausalLM,
|
||||
tokenizer: AutoTokenizer,
|
||||
input_ids: Int[Tensor, "seq"],
|
||||
choices_tokens: List[Int[Tensor, "seq"]],
|
||||
choice: Optional[Int[Tensor, ""]] = None,
|
||||
prob: float = 1,
|
||||
current_tokens: Int[Tensor, "seq"] = torch.LongTensor([]),
|
||||
z=[],
|
||||
):
|
||||
if choice is not None:
|
||||
c = choice[None].to(current_tokens.device)
|
||||
current_tokens = torch.cat([current_tokens, c], dim=-1)
|
||||
c = choice[None].to(input_ids.device)
|
||||
input_ids = torch.cat([input_ids, c], dim=-1)
|
||||
|
||||
next_choices = get_valid_next_choices(choices_tokens, current_tokens)
|
||||
if len(next_choices) == 0:
|
||||
s = tokenizer.decode(current_tokens)
|
||||
r = dict(prob=prob, choice=s)
|
||||
yield r
|
||||
else:
|
||||
o = model(input_ids[None])
|
||||
logits_constrained = o.logits[0, -1][next_choices]
|
||||
probs = F.softmax(logits_constrained, dim=-1)
|
||||
for i in range(len(next_choices)):
|
||||
next_choice = next_choices[i]
|
||||
next_prob = prob * probs[i].item()
|
||||
yield from choice_tree(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
choices_tokens=choices_tokens,
|
||||
input_ids=input_ids,
|
||||
choice=next_choice,
|
||||
prob=next_prob,
|
||||
current_tokens=current_tokens,
|
||||
z=z + [i],
|
||||
)
|
||||
@@ -5,14 +5,16 @@ from prob_jsonformer.logits_processors import (
|
||||
OutputNumbersTokens,
|
||||
StringStoppingCriteria,
|
||||
)
|
||||
from prob_jsonformer.choice_tree import choice_tree
|
||||
from termcolor import cprint
|
||||
from transformers import PreTrainedModel, PreTrainedTokenizer
|
||||
import json
|
||||
import torch
|
||||
|
||||
GENERATION_MARKER = "|GENERATION|"
|
||||
|
||||
|
||||
class prob_jsonformer:
|
||||
class Jsonformer:
|
||||
value: Dict[str, Any] = {}
|
||||
|
||||
def __init__(
|
||||
@@ -145,6 +147,18 @@ class prob_jsonformer:
|
||||
|
||||
return response.split('"')[0].strip()
|
||||
|
||||
def generate_choice_probs(self, choices) -> str:
|
||||
prompt = self.get_prompt() + '"'
|
||||
self.debug("[generate_string_prob]", prompt, is_prompt=True)
|
||||
input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(
|
||||
self.model.device
|
||||
)[0]
|
||||
choices_tokens = self.tokenizer(choices).input_ids
|
||||
choices_tokens = [torch.tensor(c) for c in choices_tokens]
|
||||
|
||||
r = list(choice_tree(self.model, self.tokenizer, input_ids, choices_tokens))
|
||||
return r # json.dumps(r)
|
||||
|
||||
def generate_object(
|
||||
self, properties: Dict[str, Any], obj: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
@@ -178,6 +192,12 @@ class prob_jsonformer:
|
||||
else:
|
||||
obj.append(self.generation_marker)
|
||||
return self.generate_string()
|
||||
elif schema_type == "choices":
|
||||
if key:
|
||||
obj[key] = self.generation_marker
|
||||
else:
|
||||
obj.append(self.generation_marker)
|
||||
return self.generate_choice_probs(schema["enum"])
|
||||
elif schema_type == "array":
|
||||
new_array = []
|
||||
obj[key] = new_array
|
||||
Reference in New Issue
Block a user