pi2-table-llama3-8b-sft_final

TabRL schema-conditioned table checkpoint

This checkpoint is the table-generation policy in the TabRL team project. It takes source text and an explicit schema and generates table rows as JSON Lines. The companion schema policy predicts the structure first, making the two stages inspectable separately.

The project uses LoRA supervised fine-tuning of Llama 3.1 8B. This repository provides a merged model and tokenizer, together with saved sample generations. It is not an RL-trained model or evidence of gains from the proposed MONA-style RL extension.

Source and contribution

• Team implementation
• Table trainer
• Table instruction builder
• Final team report
• Companion schema checkpoint

The report credits Usman Mehboob Mohammed with system design, both SFT pipelines, data preparation, training runs and integration. See the report for the full team contribution split.

Training and output format

The instruction builder reads JSONL records containing source text, a schema and a gold table. It checks schema/table column consistency and renders the target as one JSON object per row. The public trainer exposes LoRA settings and a 2K/4K/8K context schedule using VERL/FSDP and FlashAttention-2. Its source defaults should not be substituted for a missing checkpoint-specific record of final hyperparameters or data counts.

The expected output is JSON Lines with the schema's columns. This is a format target, not a guarantee that each generation is valid or factually correct.

Saved examples and evaluation

Saved table-generation samples contain source text, input schemas and generated outputs for inspection. The filename is not used here as a verified sample count. The team repository separately retains 120 schema predictions and 120 table outputs, including malformed-output records.

Read the report's comparisons with their named baselines. They do not establish state-of-the-art quality or an improvement from RL. The documentation update does not add a new benchmark or training run.

Inference example using the training prompt builder

Clone the team repository, configure compatible PyTorch, Transformers and Accelerate, and run the following from its root. The instruction builder uses only the standard library; model loading requires sufficient memory and downloads checkpoint files. The example was checked for syntax and source interfaces, not rerun with the model.

import importlib.util
import json
from pathlib import Path

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

source = Path("src/table_model/table_instructions.py")
spec = importlib.util.spec_from_file_location("table_instructions", source)
prompts = importlib.util.module_from_spec(spec)
spec.loader.exec_module(prompts)

model_id = "mohdusman001/pi2-table-llama3-8b-sft_final"
tokenizer = AutoTokenizer.from_pretrained(model_id)
if tokenizer.pad_token_id is None:
    tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
    model_id, torch_dtype=torch.bfloat16, device_map="auto"
)
model.eval()

raw_schema = {
    "table_id": "product",
    "columns": [{"name": "product"}, {"name": "price_usd"}],
}
prompt_schema = prompts.to_prompt_schema(raw_schema)
messages = [
    {"role": "system", "content": prompts.SYSTEM_PI2},
    {
        "role": "user",
        "content": prompts.build_user_message(
            "Product: Desk lamp. Price: 25 USD.", prompt_schema
        ),
    },
]
chat = tokenizer.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True
)
inputs = tokenizer(chat, return_tensors="pt").to(model.device)
with torch.inference_mode():
    output = model.generate(
        **inputs, max_new_tokens=512, do_sample=False,
        pad_token_id=tokenizer.pad_token_id,
        eos_token_id=tokenizer.eos_token_id,
    )
new_tokens = output[0, inputs["input_ids"].shape[1]:]
table_text = tokenizer.decode(new_tokens, skip_special_tokens=True)
print(table_text)

# Parsing can fail. Retain the raw output instead of silently repairing it.
try:
    rows = [json.loads(line) for line in table_text.splitlines() if line.strip()]
except json.JSONDecodeError as exc:
    raise ValueError("The model did not produce valid JSON Lines") from exc
expected_keys = [field["name"] for field in prompt_schema["fields"]]
if not rows or any(
    not isinstance(row, dict) or list(row) != expected_keys for row in rows
):
    raise ValueError("Generated rows do not match the requested schema")

A schema check does not prove that values are supported by the document. Review source fidelity separately. The small input above illustrates usage and is not a measured evaluation case.

Intended use and limits

Use this model for research on schema-conditioned generation and structured error analysis. Outputs may be malformed, incomplete or unsupported by the input. Keep the original text and raw generations for inspection, and validate before consuming generated rows. This card does not establish a production deployment or a decision-making use case.

Contact: [email protected].