Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch, json
3
4model_id = "mohdusman001/pi2-table-llama3-8b-sft_final"
5
6tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
7tokenizer.padding_side = "left"
8if tokenizer.pad_token is None:
9 tokenizer.pad_token = tokenizer.eos_token
10
11model = AutoModelForCausalLM.from_pretrained(
12 model_id,
13 torch_dtype=torch.bfloat16,
14 device_map="auto",
15 attn_implementation="flash_attention_2",
16)
17
18# Example input (text + schema)
19text = "amanda renae carraway-marsh -lrb- born january 25 , 1978 -rrb- is a beauty queen and model from manhattan , kansas who was crowned miss kansas usa 1999 . she competed in the miss usa 1999 pageant but was unplaced . she was also miss kansas teen usa 1996 ."
20schema = {
21 "table_id": "t0_kv",
22 "columns": [
23 {
24 "name": "field",
25 "path": [
26 "field"
27 ]
28 },
29 {
30 "name": "value",
31 "path": [
32 "value"
33 ]
34 }
35 ],
36 "n_cols": 2
37}
38
39user_prompt = (
40 "[SCHEMA]\n"
41 + json.dumps(schema, ensure_ascii=False)
42 + "\n<|document|>\n"
43 + text
44)
45
46messages = [{"role": "user", "content": user_prompt}]
47
48chat = tokenizer.apply_chat_template(
49 messages,
50 tokenize=False,
51 add_generation_prompt=True,
52)
53
54inputs = tokenizer(chat, return_tensors="pt").to(model.device)
55
56with torch.no_grad():
57 out = model.generate(
58 **inputs,
59 max_new_tokens=512,
60 do_sample=True,
61 temperature=0.2,
62 top_p=0.95,
63 eos_token_id=tokenizer.eos_token_id,
64 pad_token_id=tokenizer.pad_token_id,
65 )
66
67input_len = inputs["input_ids"].shape[1]
68gen_ids = out[0, input_len:]
69table_text = tokenizer.decode(gen_ids, skip_special_tokens=True)
70
71# Parse JSONL table
72rows = [json.loads(ln) for ln in table_text.splitlines() if ln.strip()]
73print(rows)infer_and_upload_pi2.py (not part of this repo by default) demonstrates
how to run batched inference on a shuffled subset of a JSONL test file, and uploads
a rich JSONL of results to the repo.pi2_table_inference_1000_samples.jsonl contains 4 important fields for each sample:input_text: the original text passage.input_schema: the JSON schema used to define the table columns.model_output: raw text that the model generated (JSONL lines).output_table: parsed JSON objects (one per row).amanda renae carraway-marsh -lrb- born january 25 , 1978 -rrb- is a beauty queen and model from manhattan , kansas who was crowned miss kansas usa 1999 . she competed in the miss usa 1999 pageant but was unplaced . she was also miss kansas teen usa 1996 .1{
2 "table_id": "t0_kv",
3 "columns": [
4 {
5 "name": "field",
6 "path": [
7 "field"
8 ]
9 },
10 {
11 "name": "value",
12 "path": [
13 "value"
14 ]
15 }
16 ],
17 "n_cols": 2
18}1[
2 "field",
3 "value"
4]
5```\n\n---\n\n## Inference (pseudo-code)
6
7```python
8from transformers import AutoModelForCausalLM, AutoTokenizer
9import torch, json
10
11model_id = "mohdusman001/pi2-table-llama3-8b-sft_final"
12
13tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
14tokenizer.padding_side = "left"
15if tokenizer.pad_token is None:
16 tokenizer.pad_token = tokenizer.eos_token
17
18model = AutoModelForCausalLM.from_pretrained(
19 model_id,
20 torch_dtype=torch.bfloat16,
21 device_map="auto",
22 attn_implementation="flash_attention_2",
23)
24
25# Example input (text + schema)
26text = "amanda renae carraway-marsh -lrb- born january 25 , 1978 -rrb- is a beauty queen and model from manhattan , kansas who was crowned miss kansas usa 1999 . she competed in the miss usa 1999 pageant but was unplaced . she was also miss kansas teen usa 1996 ."
27schema = {
28 "table_id": "t0_kv",
29 "columns": [
30 {
31 "name": "field",
32 "path": [
33 "field"
34 ]
35 },
36 {
37 "name": "value",
38 "path": [
39 "value"
40 ]
41 }
42 ],
43 "n_cols": 2
44}
45
46user_instruction = (
47 "You are a table extraction model. Given a JSON schema and a document, "
48 "you must extract rows that match the schema.\n"
49 "Return the table as JSON Lines (JSONL): one valid JSON object per line, "
50 "with keys exactly equal to the column names in the schema.\n"
51 "Do NOT output any explanations, natural language, markdown, or code — "
52 "only the JSONL table rows.\n\n"
53)
54
55user_prompt = (
56 user_instruction
57 + "[SCHEMA]\n"
58 + json.dumps(schema, ensure_ascii=False)
59 + "\n<|document|>\n"
60 + text
61)
62
63messages = [{"role": "user", "content": user_prompt}]
64
65chat = tokenizer.apply_chat_template(
66 messages,
67 tokenize=False,
68 add_generation_prompt=True,
69)
70
71inputs = tokenizer(chat, return_tensors="pt").to(model.device)
72
73with torch.no_grad():
74 out = model.generate(
75 **inputs,
76 max_new_tokens=512,
77 do_sample=False,
78 eos_token_id=tokenizer.eos_token_id,
79 pad_token_id=tokenizer.pad_token_id,
80 )
81
82input_len = inputs["input_ids"].shape[1]
83gen_ids = out[0, input_len:]
84table_text = tokenizer.decode(gen_ids, skip_special_tokens=True)
85
86# Parse JSONL table
87rows = [json.loads(ln) for ln in table_text.splitlines() if ln.strip()]
88print(rows)infer_and_upload_pi2.py demonstrates how to run batched inference on a
shuffled subset of a JSONL test file, and uploads a rich JSONL of results to the repo.pi2_table_inference_1000_samples.jsonl contains, for each sample:text: the original text passage.schema: the JSON schema used to define the table columns.model_output: raw text that the model generated (JSONL lines).rows: parsed JSON objects (one per row).table: a structured table with table_id, columns, and data matrix.amanda renae carraway-marsh -lrb- born january 25 , 1978 -rrb- is a beauty queen and model from manhattan , kansas who was crowned miss kansas usa 1999 . she competed in the miss usa 1999 pageant but was unplaced . she was also miss kansas teen usa 1996 .1{
2 "table_id": "t0_kv",
3 "columns": [
4 {
5 "name": "field",
6 "path": [
7 "field"
8 ]
9 },
10 {
11 "name": "value",
12 "path": [
13 "value"
14 ]
15 }
16 ],
17 "n_cols": 2
18}1{
2 "table_id": "t0_kv",
3 "columns": [
4 {
5 "name": "field",
6 "path": [
7 "field"
8 ]
9 },
10 {
11 "name": "value",
12 "path": [
13 "value"
14 ]
15 }
16 ],
17 "data": [
18 [
19 "amanda renae carraway-marsh",
20 "born january 25, 1978 -rrb- is a beauty queen and model from manhattan, kansas who was crowned miss kansas usa 1999. she competed in the miss usa 1999 pageant but was unplaced. she was also miss kansas teen usa 1996."
21 ]
22 ]
23}
24```\n
25
26---
27
28## Inference (pseudo-code)
29
30```python
31from transformers import AutoModelForCausalLM, AutoTokenizer
32import torch, json
33
34model_id = "mohdusman001/pi2-table-llama3-8b-sft_final"
35
36tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
37tokenizer.padding_side = "left"
38if tokenizer.pad_token is None:
39 tokenizer.pad_token = tokenizer.eos_token
40
41model = AutoModelForCausalLM.from_pretrained(
42 model_id,
43 torch_dtype=torch.bfloat16,
44 device_map="auto",
45 attn_implementation="flash_attention_2",
46)
47
48# Example input (text + schema)
49text = "amanda renae carraway-marsh -lrb- born january 25 , 1978 -rrb- is a beauty queen and model from manhattan , kansas who was crowned miss kansas usa 1999 . she competed in the miss usa 1999 pageant but was unplaced . she was also miss kansas teen usa 1996 ."
50schema = {
51 "table_id": "t0_kv",
52 "columns": [
53 {
54 "name": "field",
55 "path": [
56 "field"
57 ]
58 },
59 {
60 "name": "value",
61 "path": [
62 "value"
63 ]
64 }
65 ],
66 "n_cols": 2
67}
68
69user_instruction = (
70 "You are a table extraction model. Given a JSON schema and a document, "
71 "you must extract rows that match the schema.
72"
73 "Return the table as JSON Lines (JSONL): one valid JSON object per line, "
74 "with keys exactly equal to the column names in the schema.
75"
76 "For key-value schemas (e.g. columns ['slot', 'value'] or ['field', 'value']), "
77 "you must output one row per attribute (e.g. name, eatType, area, etc.).
78"
79 "Do NOT output any explanations, natural language, markdown, or code — "
80 "only the JSONL table rows.
81
82"
83)
84
85user_prompt = (
86 user_instruction
87 + "[SCHEMA]
88"
89 + json.dumps(schema, ensure_ascii=False)
90 + "
91<|document|>
92"
93 + text
94)
95
96messages = [{"role": "user", "content": user_prompt}]
97
98chat = tokenizer.apply_chat_template(
99 messages,
100 tokenize=False,
101 add_generation_prompt=True,
102)
103
104inputs = tokenizer(chat, return_tensors="pt").to(model.device)
105
106with torch.no_grad():
107 out = model.generate(
108 **inputs,
109 max_new_tokens=512,
110 do_sample=False,
111 eos_token_id=tokenizer.eos_token_id,
112 pad_token_id=tokenizer.pad_token_id,
113 )
114
115input_len = inputs["input_ids"].shape[1]
116gen_ids = out[0, input_len:]
117table_text = tokenizer.decode(gen_ids, skip_special_tokens=True)
118
119# Parse JSONL table
120rows = [json.loads(ln) for ln in table_text.splitlines() if ln.strip()]
121print(rows)infer_and_upload_pi2.py demonstrates how to run batched inference on a
shuffled subset of a JSONL test file, and uploads a rich JSONL of results to the repo.pi2_table_inference_1000_samples.jsonl contains, for each sample:text: the original text passage.schema: the JSON schema used to define the table columns.model_output: raw text that the model generated (JSONL lines).rows: parsed JSON objects (one per row).table: a structured table with table_id, columns, and data matrix.amanda renae carraway-marsh -lrb- born january 25 , 1978 -rrb- is a beauty queen and model from manhattan , kansas who was crowned miss kansas usa 1999 . she competed in the miss usa 1999 pageant but was unplaced . she was also miss kansas teen usa 1996 .1{
2 "table_id": "t0_kv",
3 "columns": [
4 {
5 "name": "field",
6 "path": [
7 "field"
8 ]
9 },
10 {
11 "name": "value",
12 "path": [
13 "value"
14 ]
15 }
16 ],
17 "n_cols": 2
18}1{
2 "table_id": "t0_kv",
3 "columns": [
4 {
5 "name": "field",
6 "path": [
7 "field"
8 ]
9 },
10 {
11 "name": "value",
12 "path": [
13 "value"
14 ]
15 }
16 ],
17 "data": [
18 [
19 "amanda renae carraway-marsh",
20 "born"
21 ],
22 [
23 "january 25, 1978",
24 "born"
25 ],
26 [
27 "amanda renae carraway-marsh",
28 "beauty queen and model"
29 ]
30 ]
31}