System Prompt
You are an information extraction system.
Return ONLY a valid JSON object.
Rules:
- Extract only explicitly mentioned information.
- Never infer or assume values.
- Preserve the schema exactly.
- Use null for missing values.
- Return only JSON.
User Prompt
Schema:
{
"company": "string",
"ceo": "string",
"founded": "int",
"employees": "int"
}
Text:
OpenAI was founded in 2015. Sam Altman is the CEO. The company employs approximately 4500 people.
Expected Output
{
"company": "OpenAI",
"ceo": "Sam Altman",
"founded": 2015,
"employees": 4500
}
Installation
Install the required packages.
pip install transformers peft accelerate torch
Download the Base Model
This adapter is trained on Qwen/Qwen2.5-1.5B.
The base model will be downloaded automatically the first time you run the code.
Load the Model
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
import torch
BASE_MODEL = "Qwen/Qwen2.5-1.5B"
LORA_MODEL = "Yellowforesty/qwen2.5-json-extractor-p0"
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
base = AutoModelForCausalLM.from_pretrained(
BASE_MODEL,
device_map="auto"
)
model = PeftModel.from_pretrained(
base,
LORA_MODEL
)
System Prompt
SYSTEM_PROMPT = """
You are an information extraction system.
Return ONLY a valid JSON object.
Rules:
- Extract only explicitly mentioned information.
- Never infer or assume values.
- Preserve the schema exactly.
- Use null for missing values.
- Return only JSON.
""".strip()
prompt = """
Schema:
{
"company": "string",
"ceo": "string",
"founded": "int",
"employees": "int"
}
Text:
OpenAI was founded in 2015. Sam Altman is the CEO. The company employs approximately 4500 people.
"""
Run Inference
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": prompt},
]
text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
inputs = tokenizer(text, return_tensors="pt").to(model.device)
with torch.no_grad():
output = model.generate(
**inputs,
max_new_tokens=256,
temperature=0.0,
do_sample=False,
)
response = tokenizer.decode(
output[0][inputs["input_ids"].shape[1]:],
skip_special_tokens=True,
)
print(response)
Output
{
"company": "OpenAI",
"ceo": "Sam Altman",
"founded": 2015,
"employees": 4500
}
Inference
import os
import torch
from huggingface_hub import login
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import PeftModel
BASE_MODEL = "Qwen/Qwen2.5-1.5B"
LORA_MODEL = "Yellowforesty/qwen2.5-json-extractor-p0"
HF_TOKEN = os.getenv("HF_TOKEN")
if HF_TOKEN:
login(token=HF_TOKEN)
SYSTEM_PROMPT = """
You are an information extraction system.
Return ONLY a valid JSON object.
Rules:
- Extract only explicitly mentioned information.
- Never infer or assume values.
- Preserve the schema exactly.
- Use null for missing values.
- Return only JSON.
""".strip()
PROMPT = """
Schema:
{
"company": "string",
"ceo": "string",
"founded": "int",
"employees": "int"
}
Text:
OpenAI was founded in 2015. Sam Altman is the CEO. The company employs approximately 4500 people.
""".strip()
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
base = AutoModelForCausalLM.from_pretrained(
BASE_MODEL,
device_map="auto",
torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
)
model = PeftModel.from_pretrained(
base,
LORA_MODEL,
)
model.eval()
print("Device:", next(model.parameters()).device)
messages = [
{"role": "system", "content": SYSTEM_PROMPT},
{"role": "user", "content": PROMPT},
]
text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True,
)
inputs = tokenizer(text, return_tensors="pt").to(model.device)
with torch.no_grad():
output = model.generate(
**inputs,
max_new_tokens=256,
temperature=0.0,
do_sample=False,
)
response = tokenizer.decode(
output[0][inputs["input_ids"].shape[1]:],
skip_special_tokens=True,
)
print(response.strip())
Example
Schema:
{
"name": "string",
"age": "int",
"city": "string"
}
Text:
Aman, aged 31, recently moved to Chandigarh.
Output
{
"name": "Aman",
"age": 31,
"city": "Chandigarh"
}
Benchmark
Current Version (P0)
- Primitive Types
- Flat JSON
- Null Handling
Evaluation
- JSON Validity: 100%
- Extraction Benchmark: 20 / 20
Limitations
Current version supports
- Primitive values
- Flat JSON objects
- Single object extraction
Not yet supported
- Nested JSON
- Arrays
- Number normalization
- Date normalization
- Multi-hop reasoning
- Planning
Dataset
The training dataset used to train this adapter is not included in this repository.
Roadmap
- ✅ P0 – Primitive JSON Extraction
- ⏳ P1 – Nested JSON
- ⏳ P2 – Number Normalization
- ⏳ P3 – Date Normalization
- ⏳ P4 – Reasoning
- ⏳ P5 – Planning
License
MIT