mirror of
https://github.com/wassname/prob_jsonformer.git
synced 2026-09-26 14:10:55 +08:00
better array generation
This commit is contained in:
1rgs
+62
-92
@@ -37,37 +37,42 @@
|
||||
"execution_count": 2,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
|
||||
"Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Generating...\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
|
||||
"Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Generating...\n",
|
||||
"{\n",
|
||||
" name: \u001b[32m\"Mickey Mouse\"\u001b[0m,\n",
|
||||
" age: \u001b[32m30.0\u001b[0m,\n",
|
||||
" description: \u001b[32m\"A mouse like object\"\u001b[0m\n",
|
||||
" store: {\n",
|
||||
" name: \u001b[32m\"Mike's Ski Shop\"\u001b[0m,\n",
|
||||
" location: \u001b[32m\"Sugarloaf\"\u001b[0m,\n",
|
||||
" inventory: [\n",
|
||||
" {\n",
|
||||
" productId: \u001b[32m\"1234567890\"\u001b[0m,\n",
|
||||
" name: \u001b[32m\"Ski Shop\"\u001b[0m,\n",
|
||||
" description: \u001b[32m\"Ski Shop sells premium skis and snowboards\"\u001b[0m,\n",
|
||||
" category: \u001b[32m\"Sports\"\u001b[0m,\n",
|
||||
" price: \u001b[32m20.09999\u001b[0m,\n",
|
||||
" inStock: \u001b[32mTrue\u001b[0m,\n",
|
||||
" rating: \u001b[32m5.09999\u001b[0m,\n",
|
||||
" images: [\n",
|
||||
" \u001b[32m\"https://s3.amazonaws.com/images.skisnow.com/skis\"\u001b[0m\n",
|
||||
" ]\n",
|
||||
" },\n",
|
||||
" {\n",
|
||||
" productId: \u001b[32m\"12345678910\"\u001b[0m,\n",
|
||||
" name: \u001b[32m\"Snowboard\"\u001b[0m,\n",
|
||||
" description: \u001b[32m\"Snowboard for sale\"\u001b[0m,\n",
|
||||
" category: \u001b[32m\"Sports\"\u001b[0m,\n",
|
||||
" price: \u001b[32m20.09999\u001b[0m,\n",
|
||||
" inStock: \u001b[32mTrue\u001b[0m,\n",
|
||||
" rating: \u001b[32m5.09999\u001b[0m,\n",
|
||||
" images: [\n",
|
||||
" \u001b[32m\"https://s3.amazonaws.com/images.skisnow.com/snow\"\u001b[0m\n",
|
||||
" ]\n",
|
||||
" }\n",
|
||||
" ]\n",
|
||||
" }\n",
|
||||
"}\n"
|
||||
]
|
||||
}
|
||||
@@ -76,27 +81,48 @@
|
||||
"from jsonformer.format import highlight_values\n",
|
||||
"from jsonformer.main import Jsonformer\n",
|
||||
"\n",
|
||||
"mickey_mouse_schema = {\n",
|
||||
"ecomm = {\n",
|
||||
" \"type\": \"object\",\n",
|
||||
" \"properties\": {\n",
|
||||
" \"name\": {\"type\": \"string\"},\n",
|
||||
" \"age\": {\"type\": \"number\"},\n",
|
||||
" \"description\": {\"type\": \"string\"},\n",
|
||||
" \"store\": {\n",
|
||||
" \"type\": \"object\",\n",
|
||||
" \"properties\": {\n",
|
||||
" \"name\": {\"type\": \"string\"},\n",
|
||||
" \"location\": {\"type\": \"string\"},\n",
|
||||
" \"inventory\": {\n",
|
||||
" \"type\": \"array\",\n",
|
||||
" \"items\": {\n",
|
||||
" \"type\": \"object\",\n",
|
||||
" \"properties\": {\n",
|
||||
" \"productId\": {\"type\": \"string\"},\n",
|
||||
" \"name\": {\"type\": \"string\"},\n",
|
||||
" \"description\": {\"type\": \"string\"},\n",
|
||||
" \"category\": {\"type\": \"string\"},\n",
|
||||
" \"price\": {\"type\": \"number\"},\n",
|
||||
" \"inStock\": {\"type\": \"boolean\"},\n",
|
||||
" \"rating\": {\"type\": \"number\"},\n",
|
||||
" \"images\": {\"type\": \"array\", \"items\": {\"type\": \"string\"}},\n",
|
||||
" },\n",
|
||||
" },\n",
|
||||
" },\n",
|
||||
" },\n",
|
||||
" }\n",
|
||||
" },\n",
|
||||
"}\n",
|
||||
"\n",
|
||||
"\n",
|
||||
"builder = Jsonformer(\n",
|
||||
" model=model,\n",
|
||||
" tokenizer=tokenizer,\n",
|
||||
" json_schema=mickey_mouse_schema,\n",
|
||||
" prompt=\"write a description about mickey mouse\",\n",
|
||||
" max_string_token_length=100,\n",
|
||||
" json_schema=ecomm,\n",
|
||||
" prompt=\"write a description about mike's ski shop which sells premium skis and snowboards\",\n",
|
||||
" max_string_token_length=20,\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"print(\"Generating...\")\n",
|
||||
"output = builder()\n",
|
||||
"\n",
|
||||
"highlight_values(output)\n"
|
||||
"highlight_values(output)"
|
||||
]
|
||||
},
|
||||
{
|
||||
@@ -104,35 +130,11 @@
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
|
||||
"Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n",
|
||||
"The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
|
||||
"Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Generating...\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
|
||||
"Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Generating...\n",
|
||||
"{\n",
|
||||
" make: \u001b[32m\"audi\"\u001b[0m,\n",
|
||||
" model: \u001b[32m\"a4\"\u001b[0m,\n",
|
||||
@@ -176,43 +178,11 @@
|
||||
"execution_count": 4,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
|
||||
"Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Generating...\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
|
||||
"Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n",
|
||||
"The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
|
||||
"Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n",
|
||||
"The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
|
||||
"Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n",
|
||||
"The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
|
||||
"Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n",
|
||||
"The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
|
||||
"Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n",
|
||||
"The attention mask and the pad token id were not set. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.\n",
|
||||
"Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"Generating...\n",
|
||||
"{\n",
|
||||
" car: {\n",
|
||||
" make: \u001b[32m\"Rolls Royce\"\u001b[0m,\n",
|
||||
|
||||
+17
-8
@@ -135,9 +135,7 @@ class Jsonformer:
|
||||
self.debug("[generate_string]", "|" + response + "|")
|
||||
|
||||
if response.count('"') < 1:
|
||||
raise ValueError(
|
||||
"Failed to generate a valid string, try increasing max_string_token_length"
|
||||
)
|
||||
return response
|
||||
|
||||
return response.split('"')[0].strip()
|
||||
|
||||
@@ -201,12 +199,23 @@ class Jsonformer:
|
||||
output = self.model.forward(input_tensor.to(self.model.device))
|
||||
logits = output.logits[0, -1]
|
||||
|
||||
close_bracket_token_id = self.tokenizer.convert_tokens_to_ids("]")
|
||||
comma_token_id = self.tokenizer.convert_tokens_to_ids(", ")
|
||||
close_bracket_logits = logits[close_bracket_token_id]
|
||||
comma_logits = logits[comma_token_id]
|
||||
|
||||
if close_bracket_logits > comma_logits:
|
||||
top_indices = logits.topk(30).indices
|
||||
sorted_token_ids = top_indices[logits[top_indices].argsort(descending=True)]
|
||||
|
||||
found_comma = False
|
||||
found_close_bracket = False
|
||||
|
||||
for token_id in sorted_token_ids:
|
||||
decoded_token = self.tokenizer.decode(token_id)
|
||||
if ',' in decoded_token:
|
||||
found_comma = True
|
||||
break
|
||||
if ']' in decoded_token:
|
||||
found_close_bracket = True
|
||||
break
|
||||
|
||||
if found_close_bracket or not found_comma:
|
||||
break
|
||||
|
||||
return obj
|
||||
|
||||
Reference in New Issue
Block a user