OmAlve/reading-steiner
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

初始化项目，由ModelHub XC社区提供模型

Browse Source 
Model: OmAlve/reading-steiner
Source: Original Platform
This commit is contained in:
ModelHub XC
2026-06-16 08:15:17 +08:00
commit 92eddcb2d6
18 changed files with 3106 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

36
.gitattributes vendored Normal file
View File

@@ -0,0 +1,36 @@
*.7z filter=lfs diff=lfs merge=lfs -text
*.arrow filter=lfs diff=lfs merge=lfs -text
*.bin filter=lfs diff=lfs merge=lfs -text
*.bz2 filter=lfs diff=lfs merge=lfs -text
*.ckpt filter=lfs diff=lfs merge=lfs -text
*.ftz filter=lfs diff=lfs merge=lfs -text
*.gz filter=lfs diff=lfs merge=lfs -text
*.h5 filter=lfs diff=lfs merge=lfs -text
*.joblib filter=lfs diff=lfs merge=lfs -text
*.lfs.* filter=lfs diff=lfs merge=lfs -text
*.mlmodel filter=lfs diff=lfs merge=lfs -text
*.model filter=lfs diff=lfs merge=lfs -text
*.msgpack filter=lfs diff=lfs merge=lfs -text
*.npy filter=lfs diff=lfs merge=lfs -text
*.npz filter=lfs diff=lfs merge=lfs -text
*.onnx filter=lfs diff=lfs merge=lfs -text
*.ot filter=lfs diff=lfs merge=lfs -text
*.parquet filter=lfs diff=lfs merge=lfs -text
*.pb filter=lfs diff=lfs merge=lfs -text
*.pickle filter=lfs diff=lfs merge=lfs -text
*.pkl filter=lfs diff=lfs merge=lfs -text
*.pt filter=lfs diff=lfs merge=lfs -text
*.pth filter=lfs diff=lfs merge=lfs -text
*.rar filter=lfs diff=lfs merge=lfs -text
*.safetensors filter=lfs diff=lfs merge=lfs -text
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.tar.* filter=lfs diff=lfs merge=lfs -text
*.tar filter=lfs diff=lfs merge=lfs -text
*.tflite filter=lfs diff=lfs merge=lfs -text
*.tgz filter=lfs diff=lfs merge=lfs -text
*.wasm filter=lfs diff=lfs merge=lfs -text
*.xz filter=lfs diff=lfs merge=lfs -text
*.zip filter=lfs diff=lfs merge=lfs -text
*.zst filter=lfs diff=lfs merge=lfs -text
*tfevents* filter=lfs diff=lfs merge=lfs -text
tokenizer.json filter=lfs diff=lfs merge=lfs -text

1082
HANDOFF.md Normal file
View File

File diff suppressed because it is too large Load Diff

60
README.md Normal file
View File

@@ -0,0 +1,60 @@
---
base_model: Qwen/Qwen3-0.6B
library_name: transformers
model_name: reading-steiner
tags:
- generated_from_trainer
- trackio:https://huggingface.co/spaces/OmAlve/trackio
- trl
- sft
- trackio
licence: license
---
# Model Card for reading-steiner
This model is a fine-tuned version of [Qwen/Qwen3-0.6B](https://huggingface.co/Qwen/Qwen3-0.6B).
It has been trained using [TRL](https://github.com/huggingface/trl).
## Quick start
```python
from transformers import pipeline
question = "If you had a time machine, but could only go to the past or the future once and never return, which would you choose and why?"
generator = pipeline("text-generation", model="OmAlve/reading-steiner", device="cuda")
output = generator([{"role": "user", "content": question}], max_new_tokens=128, return_full_text=False)[0]
print(output["generated_text"])
```
## Training procedure
This model was trained with SFT.
### Framework versions
- TRL: 0.24.0
- Transformers: 5.5.0
- Pytorch: 2.5.1+cu124
- Datasets: 4.3.0
- Tokenizers: 0.22.2
## Citations
Cite TRL as:
```bibtex
@misc{vonwerra2022trl,
title = {{TRL: Transformer Reinforcement Learning}},
author = {Leandro von Werra and Younes Belkada and Lewis Tunstall and Edward Beeching and Tristan Thrush and Nathan Lambert and Shengyi Huang and Kashif Rasul and Quentin Gallou{\'e}dec},
year = 2020,
journal = {GitHub repository},
publisher = {GitHub},
howpublished = {\url{https://github.com/huggingface/trl}}
}
```

89
chat_template.jinja Normal file
View File

@@ -0,0 +1,89 @@
{%- if tools %}
{{- '<|im_start|>system\n' }}
{%- if messages[0].role == 'system' %}
{{- messages[0].content + '\n\n' }}
{%- endif %}
{{- "# Tools\n\nYou may call one or more functions to assist with the user query.\n\nYou are provided with function signatures within <tools></tools> XML tags:\n<tools>" }}
{%- for tool in tools %}
{{- "\n" }}
{{- tool | tojson }}
{%- endfor %}
{{- "\n</tools>\n\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n<tool_call>\n{\"name\": <function-name>, \"arguments\": <args-json-object>}\n</tool_call><|im_end|>\n" }}
{%- else %}
{%- if messages[0].role == 'system' %}
{{- '<|im_start|>system\n' + messages[0].content + '<|im_end|>\n' }}
{%- endif %}
{%- endif %}
{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
{%- for message in messages[::-1] %}
{%- set index = (messages|length - 1) - loop.index0 %}
{%- if ns.multi_step_tool and message.role == "user" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}
{%- set ns.multi_step_tool = false %}
{%- set ns.last_query_index = index %}
{%- endif %}
{%- endfor %}
{%- for message in messages %}
{%- if message.content is string %}
{%- set content = message.content %}
{%- else %}
{%- set content = '' %}
{%- endif %}
{%- if (message.role == "user") or (message.role == "system" and not loop.first) %}
{{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
{%- elif message.role == "assistant" %}
{%- set reasoning_content = '' %}
{%- if message.reasoning_content is string %}
{%- set reasoning_content = message.reasoning_content %}
{%- else %}
{%- if '</think>' in content %}
{%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
{%- set content = content.split('</think>')[-1].lstrip('\n') %}
{%- endif %}
{%- endif %}
{%- if loop.index0 > ns.last_query_index %}
{%- if loop.last or (not loop.last and reasoning_content) %}
{{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content.strip('\n') + '\n</think>\n\n' + content.lstrip('\n') }}
{%- else %}
{{- '<|im_start|>' + message.role + '\n' + content }}
{%- endif %}
{%- else %}
{{- '<|im_start|>' + message.role + '\n' + content }}
{%- endif %}
{%- if message.tool_calls %}
{%- for tool_call in message.tool_calls %}
{%- if (loop.first and content) or (not loop.first) %}
{{- '\n' }}
{%- endif %}
{%- if tool_call.function %}
{%- set tool_call = tool_call.function %}
{%- endif %}
{{- '<tool_call>\n{"name": "' }}
{{- tool_call.name }}
{{- '", "arguments": ' }}
{%- if tool_call.arguments is string %}
{{- tool_call.arguments }}
{%- else %}
{{- tool_call.arguments | tojson }}
{%- endif %}
{{- '}\n</tool_call>' }}
{%- endfor %}
{%- endif %}
{{- '<|im_end|>\n' }}
{%- elif message.role == "tool" %}
{%- if loop.first or (messages[loop.index0 - 1].role != "tool") %}
{{- '<|im_start|>user' }}
{%- endif %}
{{- '\n<tool_response>\n' }}
{{- content }}
{{- '\n</tool_response>' }}
{%- if loop.last or (messages[loop.index0 + 1].role != "tool") %}
{{- '<|im_end|>\n' }}
{%- endif %}
{%- endif %}
{%- endfor %}
{%- if add_generation_prompt %}
{{- '<|im_start|>assistant\n' }}
{%- if enable_thinking is defined and enable_thinking is false %}
{{- '<think>\n\n</think>\n\n' }}
{%- endif %}
{%- endif %}

63
config.json Normal file
View File

@@ -0,0 +1,63 @@
{
"architectures": [
"Qwen3ForCausalLM"
],
"attention_bias": false,
"attention_dropout": 0.0,
"bos_token_id": null,
"dtype": "bfloat16",
"eos_token_id": 151645,
"head_dim": 128,
"hidden_act": "silu",
"hidden_size": 1024,
"initializer_range": 0.02,
"intermediate_size": 3072,
"layer_types": [
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention",
"full_attention"
],
"max_position_embeddings": 40960,
"max_window_layers": 28,
"model_type": "qwen3",
"num_attention_heads": 16,
"num_hidden_layers": 28,
"num_key_value_heads": 8,
"pad_token_id": 151643,
"rms_norm_eps": 1e-06,
"rope_parameters": {
"rope_theta": 1000000,
"rope_type": "default"
},
"sliding_window": null,
"tie_word_embeddings": true,
"transformers_version": "5.5.0",
"use_cache": false,
"use_sliding_window": false,
"vocab_size": 151936
}

194
eval_indexlm.py Normal file
View File

@@ -0,0 +1,194 @@
"""
Reading Steiner Evaluation Script
Tests the trained model on:
1. Query-relevant extraction (QE) - F1/Precision/Recall
2. Main content extraction (ME) - F1/Precision/Recall
3. Inference speed on CPU
"""
import json
import time
import os
import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
def parse_intervals(text):
"""Parse interval string like '[[1,3],[5,7]]' into a set of indices."""
text = text.strip()
if text.upper() == 'NA' or not text:
return set()
try:
intervals = json.loads(text)
indices = set()
for start, end in intervals:
indices.update(range(start, end + 1))
return indices
except (json.JSONDecodeError, TypeError, ValueError):
return set()
def compute_f1(pred_indices, gold_indices):
"""Compute F1, precision, recall between two sets of indices."""
if not pred_indices and not gold_indices:
return 1.0, 1.0, 1.0
if not pred_indices or not gold_indices:
return 0.0, 0.0, 0.0
tp = len(pred_indices & gold_indices)
precision = tp / len(pred_indices) if pred_indices else 0
recall = tp / len(gold_indices) if gold_indices else 0
f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0
return f1, precision, recall
def generate_response(model, tokenizer, messages, device, max_new_tokens=128):
"""Generate model response for given messages."""
text = tokenizer.apply_chat_template(
messages[:-1], # Exclude assistant message (ground truth)
tokenize=False,
add_generation_prompt=True,
enable_thinking=False,
)
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=4096).to(device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=max_new_tokens,
do_sample=False, # Greedy for deterministic eval
temperature=1.0,
pad_token_id=tokenizer.pad_token_id,
)
# Decode only the new tokens
new_tokens = outputs[0][inputs['input_ids'].shape[1]:]
response = tokenizer.decode(new_tokens, skip_special_tokens=True)
return response.strip()
def evaluate_model(model_id, device="cpu", num_samples=100):
"""Run full evaluation."""
print(f"\n{'='*60}")
print(f"Evaluating: {model_id}")
print(f"Device: {device}")
print(f"{'='*60}")
# Load model
print("Loading model...")
dtype = torch.float32 if device == "cpu" else torch.bfloat16
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=dtype,
attn_implementation="sdpa",
).to(device)
model.eval()
# Load eval dataset
print("Loading eval dataset...")
dataset = load_dataset("OmAlve/indexlm-training-data", split="eval")
# Sample
if len(dataset) > num_samples:
dataset = dataset.shuffle(seed=42).select(range(num_samples))
# Categorize examples
qe_examples = []
me_examples = []
for row in dataset:
msgs = row['messages']
system_msg = msgs[0]['content'] if msgs[0]['role'] == 'system' else ''
if 'query' in system_msg.lower() and 'relevant' in system_msg.lower():
qe_examples.append(msgs)
else:
me_examples.append(msgs)
print(f"QE examples: {len(qe_examples)}, ME examples: {len(me_examples)}")
# Evaluate QE
print("\n--- Query-Relevant Extraction (QE) ---")
qe_metrics = evaluate_task(model, tokenizer, qe_examples[:50], device)
# Evaluate ME
print("\n--- Main Content Extraction (ME) ---")
me_metrics = evaluate_task(model, tokenizer, me_examples[:50], device)
# Speed test
print("\n--- Inference Speed Test ---")
speed_test(model, tokenizer, qe_examples[:20], device)
return qe_metrics, me_metrics
def evaluate_task(model, tokenizer, examples, device):
"""Evaluate on a set of examples."""
if not examples:
print("No examples for this task.")
return {}
f1_scores = []
precision_scores = []
recall_scores = []
exact_matches = 0
for i, msgs in enumerate(examples):
gold = msgs[-1]['content']
gold_indices = parse_intervals(gold)
pred = generate_response(model, tokenizer, msgs, device)
pred_indices = parse_intervals(pred)
f1, prec, rec = compute_f1(pred_indices, gold_indices)
f1_scores.append(f1)
precision_scores.append(prec)
recall_scores.append(rec)
if pred_indices == gold_indices:
exact_matches += 1
if i < 3:
print(f" Example {i+1}:")
print(f" Gold: {gold}")
print(f" Pred: {pred}")
print(f" F1: {f1:.3f}, P: {prec:.3f}, R: {rec:.3f}")
avg_f1 = sum(f1_scores) / len(f1_scores) * 100
avg_prec = sum(precision_scores) / len(precision_scores) * 100
avg_rec = sum(recall_scores) / len(recall_scores) * 100
em_rate = exact_matches / len(examples) * 100
print(f"\n Results ({len(examples)} examples):")
print(f" F1: {avg_f1:.2f}")
print(f" Precision: {avg_prec:.2f}")
print(f" Recall: {avg_rec:.2f}")
print(f" Exact Match: {em_rate:.2f}%")
return {"f1": avg_f1, "precision": avg_prec, "recall": avg_rec, "exact_match": em_rate}
def speed_test(model, tokenizer, examples, device):
"""Test inference speed."""
if not examples:
return
times = []
for msgs in examples:
start = time.time()
_ = generate_response(model, tokenizer, msgs, device)
elapsed = time.time() - start
times.append(elapsed)
avg_time = sum(times) / len(times)
print(f" Average inference time: {avg_time:.3f}s ({device})")
print(f" Min: {min(times):.3f}s, Max: {max(times):.3f}s")
print(f" Throughput: {1/avg_time:.1f} pages/sec")
if __name__ == "__main__":
model_id = os.environ.get("MODEL_ID", "OmAlve/reading-steiner")
device = "cuda" if torch.cuda.is_available() else "cpu"
evaluate_model(model_id, device=device, num_samples=100)

12
generation_config.json Normal file
View File

@@ -0,0 +1,12 @@
{
"do_sample": true,
"eos_token_id": [
151645,
151643
],
"pad_token_id": 151643,
"temperature": 0.6,
"top_k": 20,
"top_p": 0.95,
"transformers_version": "5.5.0"
}

3
model.safetensors Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:f0b468190c0c4336d45ccab3bb36097639a63a67f58ac443efd9bd67d9d42f95
size 1192135096

486
prepare_data.py Normal file
View File

@@ -0,0 +1,486 @@
"""
Prepare IndexLM training data from HotpotQA and MSMARCO.
Pipeline:
1. Load HotpotQA (has context = list of (title, sentences) + supporting_facts)
2. Convert context into indexed HTML-like blocks: [i] <tag>content</tag>
3. The target is index intervals of blocks containing supporting facts
4. Also create main-content extraction examples (all content blocks are "main content",
but we inject noise blocks like nav/ads to train the model to filter them)
5. Format as conversational messages for SFT
"""
import json
import random
import re
from datasets import load_dataset, Dataset
from collections import defaultdict
random.seed(42)
# Noise blocks to inject (simulating real web page clutter)
NOISE_BLOCKS = [
'<nav>Home | About | Contact | Privacy Policy</nav>',
'<div class="ad">Advertisement - Continue Reading Below</div>',
'<div class="sidebar">Related Articles: Top 10 Facts You Didn\'t Know</div>',
'<footer>© 2024 All Rights Reserved | Terms of Service</footer>',
'<div class="cookie-banner">This site uses cookies. Accept | Decline</div>',
'<div class="social">Share on: Twitter | Facebook | LinkedIn</div>',
'<nav class="breadcrumb">Home > Category > Subcategory > Article</nav>',
'<div class="newsletter">Subscribe to our newsletter for updates</div>',
'<div class="popup">Sign up for free access to premium content</div>',
'<aside>Trending: Latest news and popular stories</aside>',
'<div class="comments">Comments (0) - Be the first to comment</div>',
'<div class="author">Written by Staff Reporter | Updated: Jan 2024</div>',
'<div class="pagination">Previous | 1 | 2 | 3 | Next</div>',
'<div class="search">Search this site...</div>',
'<div class="menu">Categories: Science, Tech, Health, Sports</div>',
]
SYSTEM_PROMPT_QE = """You are IndexLM, a web content extraction model. Given a webpage split into indexed blocks and a user query, identify which blocks contain content relevant to the query.
Each block is formatted as: [i] <tag>content</tag>
Output the indices of relevant blocks as a Python list of [start, end] intervals (inclusive).
If no relevant content exists, output 'NA'.
Example output: [[2,4],[7,7],[10,12]]"""
SYSTEM_PROMPT_ME = """You are IndexLM, a web content extraction model. Given a webpage split into indexed blocks, identify which blocks contain the main content of the page (filtering out navigation, advertisements, sidebars, and other non-content elements).
Each block is formatted as: [i] <tag>content</tag>
Output the indices of main content blocks as a Python list of [start, end] intervals (inclusive).
If no main content exists, output 'NA'.
Example output: [[1,3],[5,8],[11,15]]"""
def indices_to_intervals(indices):
"""Convert a sorted list of indices to intervals [[start,end], ...]"""
if not indices:
return "NA"
indices = sorted(set(indices))
intervals = []
start = indices[0]
end = indices[0]
for i in indices[1:]:
if i == end + 1:
end = i
else:
intervals.append([start, end])
start = i
end = i
intervals.append([start, end])
return json.dumps(intervals)
def create_indexed_blocks_from_hotpotqa(context, supporting_facts, inject_noise=True):
"""
Convert HotpotQA context into indexed HTML blocks.
context: {'title': [...], 'sentences': [[...], ...]}
supporting_facts: {'title': [...], 'sent_id': [...]}
Returns: (block_text, relevant_indices, all_content_indices)
"""
titles = context['title']
sentences_list = context['sentences']
# Build supporting facts lookup
sf_lookup = defaultdict(set)
for title, sent_id in zip(supporting_facts['title'], supporting_facts['sent_id']):
sf_lookup[title].add(sent_id)
blocks = []
relevant_indices = []
content_indices = [] # All real content (non-noise)
idx = 1
for doc_idx, (title, sentences) in enumerate(zip(titles, sentences_list)):
# Title block
blocks.append(f"[{idx}] <h2>{title}</h2>")
content_indices.append(idx)
if title in sf_lookup:
# Title of a supporting document is relevant
relevant_indices.append(idx)
idx += 1
# Sentence blocks
for sent_idx, sentence in enumerate(sentences):
sentence = sentence.strip()
if not sentence:
continue
# Use <p> for regular text
blocks.append(f"[{idx}] <p>{sentence}</p>")
content_indices.append(idx)
if title in sf_lookup and sent_idx in sf_lookup[title]:
relevant_indices.append(idx)
idx += 1
# Inject noise between documents sometimes
if inject_noise and random.random() < 0.4 and doc_idx < len(titles) - 1:
noise = random.choice(NOISE_BLOCKS)
blocks.append(f"[{idx}] {noise}")
idx += 1
# Sometimes add noise at start and end
if inject_noise:
prefix_noise = []
if random.random() < 0.5:
for _ in range(random.randint(1, 3)):
noise = random.choice(NOISE_BLOCKS)
prefix_noise.append(noise)
suffix_noise = []
if random.random() < 0.5:
for _ in range(random.randint(1, 3)):
noise = random.choice(NOISE_BLOCKS)
suffix_noise.append(noise)
if prefix_noise or suffix_noise:
# Reindex everything
new_blocks = []
new_relevant = []
new_content = []
new_idx = 1
# Prefix noise
for noise in prefix_noise:
new_blocks.append(f"[{new_idx}] {noise}")
new_idx += 1
# Remap original blocks
offset = len(prefix_noise)
for b in blocks:
old_idx = int(b.split(']')[0].replace('[', ''))
new_b = f"[{old_idx + offset}] " + '] '.join(b.split('] ')[1:])
new_blocks.append(new_b)
new_relevant = [r + offset for r in relevant_indices]
new_content = [c + offset for c in content_indices]
# Suffix noise
next_idx = len(new_blocks) + 1
for noise in suffix_noise:
new_blocks.append(f"[{next_idx}] {noise}")
next_idx += 1
blocks = new_blocks
relevant_indices = new_relevant
content_indices = new_content
block_text = "\n".join(blocks)
return block_text, relevant_indices, content_indices
def build_query_relevant_example(question, block_text, relevant_indices, url="https://en.wikipedia.org"):
"""Build a query-relevant extraction (QE) example."""
intervals = indices_to_intervals(relevant_indices)
user_content = f"URL: {url}\nQuery: {question}\n\nBlocks:\n{block_text}\n\nOutput the index intervals of blocks relevant to the query."
messages = [
{"role": "system", "content": SYSTEM_PROMPT_QE},
{"role": "user", "content": user_content},
{"role": "assistant", "content": intervals}
]
return messages
def build_main_content_example(block_text, content_indices, title="Wikipedia Article", url="https://en.wikipedia.org"):
"""Build a main content extraction (ME) example."""
intervals = indices_to_intervals(content_indices)
user_content = f"URL: {url}\nTitle: {title}\n\nBlocks:\n{block_text}\n\nOutput the index intervals of main content blocks."
messages = [
{"role": "system", "content": SYSTEM_PROMPT_ME},
{"role": "user", "content": user_content},
{"role": "assistant", "content": intervals}
]
return messages
def process_hotpotqa():
"""Process HotpotQA into IndexLM training data."""
print("Loading HotpotQA...")
ds = load_dataset("hotpotqa/hotpot_qa", "distractor", split="train")
# Sample a manageable amount
num_samples = min(15000, len(ds))
ds = ds.shuffle(seed=42).select(range(num_samples))
all_examples = []
skipped = 0
for i, row in enumerate(ds):
if i % 1000 == 0:
print(f"Processing {i}/{num_samples}...")
try:
block_text, relevant_indices, content_indices = create_indexed_blocks_from_hotpotqa(
row['context'], row['supporting_facts'], inject_noise=True
)
# Skip if too few relevant indices
if len(relevant_indices) < 1:
skipped += 1
continue
# Query-relevant extraction example
qe_messages = build_query_relevant_example(
row['question'], block_text, relevant_indices
)
all_examples.append({
"messages": qe_messages,
"task_type": "query_relevant",
"source": "hotpotqa"
})
# Main content extraction example (50% of the time)
if random.random() < 0.5:
me_messages = build_main_content_example(
block_text, content_indices,
title=row['context']['title'][0] if row['context']['title'] else "Article"
)
all_examples.append({
"messages": me_messages,
"task_type": "main_content",
"source": "hotpotqa"
})
except Exception as e:
skipped += 1
if skipped < 5:
print(f"Error on row {i}: {e}")
continue
print(f"Created {len(all_examples)} examples from HotpotQA ({skipped} skipped)")
return all_examples
def create_synthetic_web_pages():
"""Create synthetic web page examples for main content extraction training."""
print("Creating synthetic web page examples...")
# Load a text dataset to get content
ds = load_dataset("hotpotqa/hotpot_qa", "distractor", split="validation")
ds = ds.shuffle(seed=123).select(range(3000))
examples = []
for i, row in enumerate(ds):
if i % 500 == 0:
print(f"Synthetic page {i}/3000...")
try:
# Build a more realistic web page structure
titles = row['context']['title']
sentences_list = row['context']['sentences']
if not titles or not sentences_list:
continue
blocks = []
content_indices = []
idx = 1
# Header noise (nav, etc.)
num_header_noise = random.randint(1, 4)
for _ in range(num_header_noise):
blocks.append(f"[{idx}] {random.choice(NOISE_BLOCKS)}")
idx += 1
# Page title
main_title = titles[0]
blocks.append(f"[{idx}] <h1>{main_title}</h1>")
content_indices.append(idx)
idx += 1
# Main content (just first 1-3 documents)
num_docs = min(random.randint(1, 3), len(titles))
for doc_idx in range(num_docs):
title = titles[doc_idx]
sents = sentences_list[doc_idx]
if doc_idx > 0:
blocks.append(f"[{idx}] <h2>{title}</h2>")
content_indices.append(idx)
idx += 1
for sent in sents:
sent = sent.strip()
if not sent:
continue
blocks.append(f"[{idx}] <p>{sent}</p>")
content_indices.append(idx)
idx += 1
# Occasional inline noise
if random.random() < 0.3:
blocks.append(f"[{idx}] {random.choice(NOISE_BLOCKS)}")
idx += 1
# Footer noise
num_footer_noise = random.randint(1, 4)
for _ in range(num_footer_noise):
blocks.append(f"[{idx}] {random.choice(NOISE_BLOCKS)}")
idx += 1
block_text = "\n".join(blocks)
me_messages = build_main_content_example(
block_text, content_indices,
title=main_title,
url=f"https://en.wikipedia.org/wiki/{main_title.replace(' ', '_')}"
)
examples.append({
"messages": me_messages,
"task_type": "main_content",
"source": "synthetic"
})
except Exception as e:
continue
print(f"Created {len(examples)} synthetic web page examples")
return examples
def create_na_examples():
"""Create examples where no relevant content exists (model should output 'NA')."""
print("Creating NA examples...")
ds = load_dataset("hotpotqa/hotpot_qa", "distractor", split="validation")
ds = ds.shuffle(seed=456).select(range(1000))
examples = []
for i, row in enumerate(ds):
try:
# Use context from one question but query from another (mismatched)
other_idx = (i + 500) % len(ds)
other_question = ds[other_idx]['question']
# Build blocks from current context but keep only non-supporting content
block_text, _, content_indices = create_indexed_blocks_from_hotpotqa(
row['context'], {'title': [], 'sent_id': []}, inject_noise=True
)
# The query doesn't match this content → expected output: NA
# But actually some content might still be tangentially relevant,
# so we'll be conservative and only do this for clearly mismatched pairs
user_content = f"URL: https://en.wikipedia.org\nQuery: {other_question}\n\nBlocks:\n{block_text}\n\nOutput the index intervals of blocks relevant to the query."
messages = [
{"role": "system", "content": SYSTEM_PROMPT_QE},
{"role": "user", "content": user_content},
{"role": "assistant", "content": "NA"}
]
examples.append({
"messages": messages,
"task_type": "query_relevant_na",
"source": "hotpotqa_mismatched"
})
except:
continue
# Keep only a fraction (the paper mentions partial filtering of NA)
random.shuffle(examples)
examples = examples[:300]
print(f"Created {len(examples)} NA examples")
return examples
def main():
# Build all training examples
qe_examples = process_hotpotqa()
me_examples = create_synthetic_web_pages()
na_examples = create_na_examples()
all_examples = qe_examples + me_examples + na_examples
random.shuffle(all_examples)
print(f"\nTotal examples: {len(all_examples)}")
# Count by type
type_counts = defaultdict(int)
for ex in all_examples:
type_counts[ex['task_type']] += 1
for t, c in type_counts.items():
print(f" {t}: {c}")
# Check lengths
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
lengths = []
for ex in all_examples[:500]:
text = tokenizer.apply_chat_template(ex['messages'], tokenize=False)
tokens = tokenizer.encode(text)
lengths.append(len(tokens))
print(f"\nToken length stats (sample of 500):")
print(f" Min: {min(lengths)}")
print(f" Max: {max(lengths)}")
print(f" Mean: {sum(lengths)/len(lengths):.0f}")
print(f" Median: {sorted(lengths)[len(lengths)//2]}")
# Filter out examples that are too long (>4096 tokens for efficiency)
MAX_LEN = 4096
filtered = []
too_long = 0
for ex in all_examples:
text = tokenizer.apply_chat_template(ex['messages'], tokenize=False)
tokens = tokenizer.encode(text)
if len(tokens) <= MAX_LEN:
filtered.append(ex)
else:
too_long += 1
print(f"\nFiltered: {too_long} examples too long (>{MAX_LEN} tokens)")
print(f"Final dataset: {len(filtered)} examples")
# Split into train/eval
random.shuffle(filtered)
eval_size = min(500, len(filtered) // 10)
train_data = filtered[:-eval_size]
eval_data = filtered[-eval_size:]
print(f"Train: {len(train_data)}, Eval: {len(eval_data)}")
# Create HF dataset with just messages column (for SFTTrainer)
train_ds = Dataset.from_list([{"messages": ex["messages"]} for ex in train_data])
eval_ds = Dataset.from_list([{"messages": ex["messages"]} for ex in eval_data])
# Save locally
train_ds.save_to_disk("/app/indexlm_train")
eval_ds.save_to_disk("/app/indexlm_eval")
# Also push to HF Hub
from huggingface_hub import login
import os
login(token=os.environ.get("HF_TOKEN"))
from datasets import DatasetDict
ds_dict = DatasetDict({"train": train_ds, "eval": eval_ds})
ds_dict.push_to_hub("OmAlve/indexlm-training-data")
print("\nDone! Dataset pushed to OmAlve/indexlm-training-data")
# Print sample
print("\n=== Sample QE example ===")
for ex in train_data[:3]:
if ex.get("task_type", "") == "query_relevant":
for m in ex["messages"]:
print(f"\n[{m['role']}]: {m['content'][:200]}...")
break
print("\n=== Sample ME example ===")
for ex in train_data[:10]:
if ex.get("task_type", "") == "main_content":
for m in ex["messages"]:
print(f"\n[{m['role']}]: {m['content'][:200]}...")
break
if __name__ == "__main__":
main()

897
prepare_data_v2.py Normal file
View File

@@ -0,0 +1,897 @@
"""
Prepare DIVERSE IndexLM training data from multiple sources:
1. HtmlRAG-train (real Bing-scraped web HTML) — diverse domains
2. MultiHopRAG (news domain) — technology, business, sports, entertainment
3. HotpotQA (Wikipedia) — structured QA with supporting facts
This avoids the Wikipedia-only bias of the original dataset.
Output: Conversational messages for SFT with TRL SFTTrainer
Format: system + user (indexed HTML blocks + query) → assistant (index intervals)
"""
import json
import random
import re
import os
from datasets import load_dataset, Dataset, DatasetDict
from collections import defaultdict
from bs4 import BeautifulSoup
import html as html_lib
random.seed(42)
# ============ System Prompts ============
SYSTEM_PROMPT_QE = """You are IndexLM, a web content extraction model. Given a webpage split into indexed blocks and a user query, identify which blocks contain content relevant to the query.
Each block is formatted as: [i] <tag>content</tag>
Output the indices of relevant blocks as a Python list of [start, end] intervals (inclusive).
If no relevant content exists, output 'NA'.
Example output: [[2,4],[7,7],[10,12]]"""
SYSTEM_PROMPT_ME = """You are IndexLM, a web content extraction model. Given a webpage split into indexed blocks, identify which blocks contain the main content of the page (filtering out navigation, advertisements, sidebars, and other non-content elements).
Each block is formatted as: [i] <tag>content</tag>
Output the indices of main content blocks as a Python list of [start, end] intervals (inclusive).
If no main content exists, output 'NA'.
Example output: [[1,3],[5,8],[11,15]]"""
# ============ Noise blocks for injection ============
NOISE_BLOCKS_REALISTIC = [
'<nav>Home | About | Contact | Privacy Policy | Terms of Service</nav>',
'<div class="ad">Advertisement - Continue Reading Below</div>',
'<div class="sidebar">Related Articles: Top 10 Facts You Didn\'t Know</div>',
'<footer>© 2024 All Rights Reserved | Terms of Service | Cookie Policy</footer>',
'<div class="cookie-banner">This website uses cookies to improve your experience. By continuing to use this site, you consent to our use of cookies. Accept | Manage Preferences</div>',
'<div class="social-share">Share: <a>Twitter</a> | <a>Facebook</a> | <a>LinkedIn</a> | <a>Reddit</a> | <a>Email</a></div>',
'<nav class="breadcrumb">Home > Category > Subcategory > Current Article</nav>',
'<div class="newsletter-signup">Subscribe to our newsletter for the latest updates delivered to your inbox weekly.</div>',
'<div class="popup-overlay">Sign up for free access to premium content! Enter your email below.</div>',
'<aside class="trending">Trending Now: Latest breaking news and popular stories from around the web</aside>',
'<div class="comments-section">Comments (0) — Be the first to comment! Please read our community guidelines before posting.</div>',
'<div class="author-bio">Written by Staff Reporter | Updated: January 15, 2024 | 5 min read</div>',
'<div class="pagination">← Previous Article | Page 1 of 3 | Next Article →</div>',
'<div class="search-bar"><form>Search this site... <button>Go</button></form></div>',
'<div class="category-menu">Categories: Science | Technology | Health | Business | Sports | Entertainment | Politics</div>',
'<div class="login-prompt">Already a subscriber? Log in for full access. Not a member? Subscribe now starting at $4.99/month.</div>',
'<div class="related-articles"><h3>You May Also Like</h3><ul><li>10 Things You Didn\'t Know About...</li><li>Breaking: Latest Update on...</li></ul></div>',
'<div class="video-embed">Watch: Video player requires JavaScript to be enabled. [Video placeholder]</div>',
'<div class="breaking-news-ticker">BREAKING: Markets rally on latest economic data | Sports: Championship results | Weather: Storm warning issued</div>',
'<div class="accessibility">Skip to main content | Skip to navigation | Accessibility statement</div>',
'<div class="gdpr-notice">We value your privacy. We and our partners use tracking technologies to improve your browsing experience, serve personalized content, and analyze traffic.</div>',
'<div class="app-download">Download our app for a better reading experience! Available on iOS and Android.</div>',
'<script>/* Google Analytics tracking code */</script>',
'<div class="print-notice">This article is available in print edition. Subscribe for home delivery.</div>',
'<div class="sponsored">Sponsored Content | Advertiser Disclosure: Some links on this page are affiliate links.</div>',
'<div class="feedback">Was this article helpful? Yes | No | Send Feedback</div>',
'<div class="language-selector">Language: English | Español | Français | Deutsch | 日本語 | 中文</div>',
'<div class="site-footer"><ul><li>About Us</li><li>Careers</li><li>Advertise</li><li>Press</li><li>Help Center</li><li>Sitemap</li></ul></div>',
]
def indices_to_intervals(indices):
"""Convert a sorted list of indices to intervals [[start,end], ...]"""
if not indices:
return "NA"
indices = sorted(set(indices))
intervals = []
start = indices[0]
end = indices[0]
for i in indices[1:]:
if i == end + 1:
end = i
else:
intervals.append([start, end])
start = i
end = i
intervals.append([start, end])
return json.dumps(intervals)
# ============================================================
# SOURCE 1: HtmlRAG-train (Real Bing-scraped web HTML)
# ============================================================
def extract_text_content(html_str):
"""Extract visible text from an HTML string."""
try:
soup = BeautifulSoup(html_str, 'html.parser')
return soup.get_text(separator=' ', strip=True)
except:
# Fallback: strip tags with regex
clean = re.sub(r'<[^>]+>', ' ', html_str)
return re.sub(r'\s+', ' ', clean).strip()
def segment_html_to_blocks(html_content):
"""
Segment real HTML content into indexed blocks.
Splits by block-level HTML tags and line boundaries.
"""
blocks = []
# Strategy: split by block-level closing/opening tags
# HtmlRAG uses tags like <div0>, <p>, <h20>, <li>, etc.
# Split at positions where block-level tags start
block_tag_pattern = r'(<(?:div|p|h[1-6]|li|ul|ol|table|tr|td|th|article|section|header|footer|nav|aside|main|blockquote|pre|form|figure|figcaption|details|summary|option|title|button|label|select|textarea|hgroup|dl|dd|dt|caption|thead|tbody|tfoot)\b[^>]*>)'
# Also handle HtmlRAG numbered tags like <div0>, <h20>, etc.
block_tag_pattern_numbered = r'(<(?:div|p|h|li|ul|ol|table|tr|td|th|article|section|header|footer|nav|aside|main|blockquote|pre|form|figure|option|title|button|hgroup)\d*[^>]*>)'
# Split content by block-level tags
parts = re.split(block_tag_pattern_numbered, html_content)
current_block = ''
for part in parts:
part = part.strip()
if not part:
continue
# Check if this part is a block-level opening tag
if re.match(block_tag_pattern_numbered, part):
# Save previous block if it has content
if current_block.strip():
blocks.append(current_block.strip())
current_block = part
else:
current_block += ' ' + part
# Don't forget the last block
if current_block.strip():
blocks.append(current_block.strip())
# If tag-based splitting yields too few blocks, fall back to line-based
if len(blocks) < 5:
blocks = []
lines = html_content.split('\n')
for line in lines:
line = line.strip()
if line and len(line) > 5:
blocks.append(line)
# If still too few, split by multiple tags on same line
if len(blocks) < 5:
new_blocks = []
for block in blocks:
# Try splitting long blocks by inner tags
if len(block) > 200:
inner_parts = re.split(r'(</(?:div|p|h[1-6]|li|td|th|article|section)\d*>)', block)
current = ''
for ip in inner_parts:
current += ip
if re.match(r'</(?:div|p|h[1-6]|li|td|th|article|section)\d*>', ip):
if current.strip():
new_blocks.append(current.strip())
current = ''
if current.strip():
new_blocks.append(current.strip())
else:
new_blocks.append(block)
if len(new_blocks) > len(blocks):
blocks = new_blocks
# Filter: extract text and remove blocks with no meaningful content
def extract_text_simple(s):
clean = re.sub(r'<[^>]+>', ' ', s)
return re.sub(r'\s+', ' ', clean).strip()
blocks = [b for b in blocks if len(extract_text_simple(b)) > 5]
return blocks
def classify_block_as_noise(block_text):
"""Heuristic: classify if a block is likely noise (nav, ad, etc.)."""
text_lower = block_text.lower()
noise_indicators = [
'cookie', 'privacy policy', 'terms of service', 'advertisement',
'subscribe', 'newsletter', 'sign up', 'log in', 'login',
'copyright ©', 'all rights reserved', 'skip to', 'accessibility',
'share on twitter', 'share on facebook', 'social media',
'related articles', 'you may also like', 'trending now',
'app download', 'sponsored content', 'affiliate',
]
nav_patterns = ['<nav', '<footer', '<aside', 'class="ad"', 'class="sidebar"',
'class="menu"', 'class="social"', 'class="cookie"']
for indicator in noise_indicators:
if indicator in text_lower:
return True
for pattern in nav_patterns:
if pattern in text_lower:
return True
return False
def process_htmlrag_example(row):
"""Convert an HtmlRAG example to IndexLM format."""
user_content = row['messages'][0]['content']
assistant_content = row['messages'][1]['content']
score = row.get('score', 0)
# Skip low-quality examples
if score < 0.5:
return None
# Parse out HTML and question
parts = user_content.split('**Question**:')
if len(parts) < 2:
parts = user_content.split('**Question**')
if len(parts) < 2:
return None
html_raw = parts[0]
question_raw = parts[1].strip()
# Clean up the HTML marker
html_raw = html_raw.replace('**HTML**: ```', '').rstrip('`').strip()
# Extract just the question (remove the instruction part)
question = question_raw.split('\n')[0].strip().strip('*').strip()
if not question:
return None
# Segment HTML into blocks
blocks = segment_html_to_blocks(html_raw)
if len(blocks) < 3:
return None
# Get the relevant content from assistant output
relevant_text = extract_text_content(assistant_content)
relevant_words = set(relevant_text.lower().split())
# Build indexed blocks and find relevant ones
indexed_blocks = []
relevant_indices = []
content_indices = []
for idx, block in enumerate(blocks, 1):
# Determine the best tag for this block
tag_match = re.match(r'<(\w+)', block)
if tag_match:
tag = tag_match.group(1)
# Normalize numbered tags (div0 -> div, h20 -> h2)
tag = re.sub(r'\d+$', '', tag)
if not tag:
tag = 'div'
else:
tag = 'p'
text = extract_text_content(block)
if not text or len(text) < 3:
continue
indexed_blocks.append(f"[{idx}] <{tag}>{text}</{tag}>")
# Check if this block is noise
is_noise = classify_block_as_noise(block)
if not is_noise:
content_indices.append(idx)
# Check relevance by substring matching with assistant output
# Use the full relevant text as a search target
text_lower = text.lower()
relevant_lower = relevant_text.lower()
# Method 1: Check if significant portions of relevant text appear in block
# Split relevant text into 3-word ngrams and check for matches
rel_words_list = relevant_lower.split()
matched = False
# Check 3-gram overlap
for i in range(len(rel_words_list) - 2):
trigram = ' '.join(rel_words_list[i:i+3])
if trigram in text_lower:
matched = True
break
# Also check: does the block text appear as a substring in the relevant text?
if not matched and len(text) > 15:
# Check if meaningful portion of block appears in relevant output
block_sentences = [s.strip() for s in text.split('.') if len(s.strip()) > 10]
for sent in block_sentences:
if sent.lower() in relevant_lower:
matched = True
break
# Also check word overlap with a more lenient threshold
if not matched:
block_words = set(text_lower.split())
if relevant_words and block_words:
overlap_count = len(block_words & relevant_words)
# At least 3 content words overlap (excluding stopwords)
stopwords = {'the','a','an','is','are','was','were','in','on','at','to','for','of','and','or','but','with','by','from','as','it','this','that','be','has','have','had','do','does','did','not','no'}
content_overlap = len((block_words - stopwords) & (relevant_words - stopwords))
if content_overlap >= 2:
matched = True
if matched:
relevant_indices.append(idx)
if not indexed_blocks or len(indexed_blocks) < 3:
return None
block_text = "\n".join(indexed_blocks)
results = []
# Query-relevant extraction example
if relevant_indices:
intervals = indices_to_intervals(relevant_indices)
user_msg = f"URL: https://example.com\nQuery: {question}\n\nBlocks:\n{block_text}\n\nOutput the index intervals of blocks relevant to the query."
results.append({
"messages": [
{"role": "system", "content": SYSTEM_PROMPT_QE},
{"role": "user", "content": user_msg},
{"role": "assistant", "content": intervals}
],
"task_type": "query_relevant",
"source": "htmlrag"
})
# Main content extraction example (30% of the time to balance)
if content_indices and random.random() < 0.3:
intervals = indices_to_intervals(content_indices)
user_msg = f"URL: https://example.com\nTitle: Web Page\n\nBlocks:\n{block_text}\n\nOutput the index intervals of main content blocks."
results.append({
"messages": [
{"role": "system", "content": SYSTEM_PROMPT_ME},
{"role": "user", "content": user_msg},
{"role": "assistant", "content": intervals}
],
"task_type": "main_content",
"source": "htmlrag"
})
return results
def load_htmlrag_data():
"""Load and convert HtmlRAG-train data."""
print("Loading HtmlRAG-train (real web HTML)...")
# Use 4k and 8k token variants - good balance of context
files = [
'nq-4k.jsonl', 'nq-8k.jsonl',
'asqa-4k.jsonl', 'asqa-8k.jsonl',
'trivia-qa-4k.jsonl', 'trivia-qa-8k.jsonl',
'musique-4k.jsonl', 'musique-8k.jsonl',
'hotpot-qa-4k.jsonl', 'hotpot-qa-8k.jsonl',
]
all_examples = []
for file in files:
print(f" Processing {file}...")
try:
ds = load_dataset('zstanjj/HtmlRAG-train', data_files=file, split='train')
count = 0
for row in ds:
results = process_htmlrag_example(row)
if results:
all_examples.extend(results)
count += len(results)
print(f" Got {count} examples from {file}")
except Exception as e:
print(f" Error loading {file}: {e}")
print(f" Total HtmlRAG examples: {len(all_examples)}")
return all_examples
# ============================================================
# SOURCE 2: MultiHopRAG (News domain)
# ============================================================
def process_multihoprag():
"""Convert MultiHopRAG news articles into IndexLM format."""
print("Loading MultiHopRAG (news domain)...")
corpus = load_dataset("yixuantt/MultiHopRAG", name="corpus", split="train")
queries = load_dataset("yixuantt/MultiHopRAG", name="MultiHopRAG", split="train")
# Build URL->article lookup
url_to_article = {}
for article in corpus:
url_to_article[article['url']] = article
all_examples = []
for q_row in queries:
query = q_row['query']
evidence_list = q_row['evidence_list']
for evidence in evidence_list:
url = evidence.get('url', '')
fact = evidence.get('fact', '')
if url not in url_to_article or not fact:
continue
article = url_to_article[url]
title = article.get('title', 'News Article')
body = article.get('body', '')
source = article.get('source', 'Unknown')
category = article.get('category', 'general')
if not body or len(body) < 100:
continue
# Split article body into paragraphs
paragraphs = [p.strip() for p in body.split('\n') if p.strip() and len(p.strip()) > 20]
if not paragraphs:
continue
# Build indexed blocks with realistic web structure
blocks = []
content_indices = []
relevant_indices = []
idx = 1
# Add realistic header noise
num_header = random.randint(1, 3)
for _ in range(num_header):
blocks.append(f"[{idx}] {random.choice(NOISE_BLOCKS_REALISTIC)}")
idx += 1
# Article title
blocks.append(f"[{idx}] <h1>{title}</h1>")
content_indices.append(idx)
idx += 1
# Author/date line
author = article.get('author', 'Staff Writer')
published = article.get('published_at', '2024-01-01')
blocks.append(f"[{idx}] <div class=\"byline\">By {author} | {source} | {published} | Category: {category}</div>")
content_indices.append(idx)
idx += 1
# Article paragraphs
fact_words = set(fact.lower().split())
for para in paragraphs:
# Determine tag
if len(para) < 60 and not para.endswith('.'):
tag = 'h2'
else:
tag = 'p'
blocks.append(f"[{idx}] <{tag}>{para}</{tag}>")
content_indices.append(idx)
# Check if paragraph contains the evidence fact
para_words = set(para.lower().split())
overlap = len(para_words & fact_words)
if overlap > 5 or (fact_words and overlap / len(fact_words) > 0.3):
relevant_indices.append(idx)
idx += 1
# Occasional mid-article noise
if random.random() < 0.15:
blocks.append(f"[{idx}] {random.choice(NOISE_BLOCKS_REALISTIC)}")
idx += 1
# Footer noise
num_footer = random.randint(1, 4)
for _ in range(num_footer):
blocks.append(f"[{idx}] {random.choice(NOISE_BLOCKS_REALISTIC)}")
idx += 1
block_text = "\n".join(blocks)
# Query-relevant extraction
if relevant_indices:
intervals = indices_to_intervals(relevant_indices)
user_msg = f"URL: {url}\nQuery: {query}\n\nBlocks:\n{block_text}\n\nOutput the index intervals of blocks relevant to the query."
all_examples.append({
"messages": [
{"role": "system", "content": SYSTEM_PROMPT_QE},
{"role": "user", "content": user_msg},
{"role": "assistant", "content": intervals}
],
"task_type": "query_relevant",
"source": "multihoprag_news"
})
# Main content extraction
if content_indices and random.random() < 0.4:
intervals = indices_to_intervals(content_indices)
user_msg = f"URL: {url}\nTitle: {title}\n\nBlocks:\n{block_text}\n\nOutput the index intervals of main content blocks."
all_examples.append({
"messages": [
{"role": "system", "content": SYSTEM_PROMPT_ME},
{"role": "user", "content": user_msg},
{"role": "assistant", "content": intervals}
],
"task_type": "main_content",
"source": "multihoprag_news"
})
print(f" Total MultiHopRAG examples: {len(all_examples)}")
return all_examples
# ============================================================
# SOURCE 3: HotpotQA (Wikipedia - but balanced as minority)
# ============================================================
def process_hotpotqa():
"""Process HotpotQA — kept but as a smaller proportion."""
print("Loading HotpotQA (Wikipedia domain)...")
ds = load_dataset("hotpotqa/hotpot_qa", "distractor", split="train")
# Reduced from 15K to 5K — wiki should be minority source
num_samples = min(5000, len(ds))
ds = ds.shuffle(seed=42).select(range(num_samples))
all_examples = []
skipped = 0
for i, row in enumerate(ds):
if i % 1000 == 0:
print(f" Processing {i}/{num_samples}...")
try:
titles = row['context']['title']
sentences_list = row['context']['sentences']
sf = row['supporting_facts']
sf_lookup = defaultdict(set)
for title, sent_id in zip(sf['title'], sf['sent_id']):
sf_lookup[title].add(sent_id)
blocks = []
relevant_indices = []
content_indices = []
idx = 1
# Header noise
if random.random() < 0.6:
for _ in range(random.randint(1, 3)):
blocks.append(f"[{idx}] {random.choice(NOISE_BLOCKS_REALISTIC)}")
idx += 1
for doc_idx, (title, sentences) in enumerate(zip(titles, sentences_list)):
blocks.append(f"[{idx}] <h2>{title}</h2>")
content_indices.append(idx)
if title in sf_lookup:
relevant_indices.append(idx)
idx += 1
for sent_idx, sentence in enumerate(sentences):
sentence = sentence.strip()
if not sentence:
continue
blocks.append(f"[{idx}] <p>{sentence}</p>")
content_indices.append(idx)
if title in sf_lookup and sent_idx in sf_lookup[title]:
relevant_indices.append(idx)
idx += 1
if random.random() < 0.3 and doc_idx < len(titles) - 1:
blocks.append(f"[{idx}] {random.choice(NOISE_BLOCKS_REALISTIC)}")
idx += 1
# Footer noise
if random.random() < 0.6:
for _ in range(random.randint(1, 3)):
blocks.append(f"[{idx}] {random.choice(NOISE_BLOCKS_REALISTIC)}")
idx += 1
if len(relevant_indices) < 1:
skipped += 1
continue
block_text = "\n".join(blocks)
# QE example
intervals = indices_to_intervals(relevant_indices)
user_msg = f"URL: https://en.wikipedia.org\nQuery: {row['question']}\n\nBlocks:\n{block_text}\n\nOutput the index intervals of blocks relevant to the query."
all_examples.append({
"messages": [
{"role": "system", "content": SYSTEM_PROMPT_QE},
{"role": "user", "content": user_msg},
{"role": "assistant", "content": intervals}
],
"task_type": "query_relevant",
"source": "hotpotqa_wiki"
})
# ME example (less frequent - wiki is minority)
if random.random() < 0.3:
intervals = indices_to_intervals(content_indices)
user_msg = f"URL: https://en.wikipedia.org\nTitle: {titles[0]}\n\nBlocks:\n{block_text}\n\nOutput the index intervals of main content blocks."
all_examples.append({
"messages": [
{"role": "system", "content": SYSTEM_PROMPT_ME},
{"role": "user", "content": user_msg},
{"role": "assistant", "content": intervals}
],
"task_type": "main_content",
"source": "hotpotqa_wiki"
})
except Exception as e:
skipped += 1
continue
print(f" Total HotpotQA examples: {len(all_examples)} ({skipped} skipped)")
return all_examples
# ============================================================
# SOURCE 4: MS MARCO (Diverse web QA)
# ============================================================
def process_msmarco():
"""Process MS MARCO for diverse web domain QA examples."""
print("Loading MS MARCO (diverse web QA)...")
try:
ds = load_dataset("microsoft/ms_marco", "v1.1", split="train")
# Sample a manageable subset
num_samples = min(5000, len(ds))
ds = ds.shuffle(seed=99).select(range(num_samples))
except Exception as e:
print(f" Could not load MS MARCO: {e}")
return []
all_examples = []
for i, row in enumerate(ds):
if i % 1000 == 0:
print(f" Processing {i}/{num_samples}...")
try:
query = row['query']
passages = row['passages']
if not passages or not passages.get('passage_text'):
continue
passage_texts = passages['passage_text']
is_selected = passages.get('is_selected', [0] * len(passage_texts))
if not any(is_selected):
continue
# Build blocks from passages (these are real web snippets from Bing)
blocks = []
relevant_indices = []
content_indices = []
idx = 1
# Header noise
if random.random() < 0.5:
for _ in range(random.randint(1, 2)):
blocks.append(f"[{idx}] {random.choice(NOISE_BLOCKS_REALISTIC)}")
idx += 1
for p_idx, (text, selected) in enumerate(zip(passage_texts, is_selected)):
text = text.strip()
if not text:
continue
# Simulate different content types
if p_idx == 0 and random.random() < 0.3:
tag = 'h1'
elif len(text) < 80:
tag = random.choice(['h2', 'h3', 'strong'])
else:
tag = 'p'
blocks.append(f"[{idx}] <{tag}>{text}</{tag}>")
content_indices.append(idx)
if selected:
relevant_indices.append(idx)
idx += 1
# Between-passage noise
if random.random() < 0.2:
blocks.append(f"[{idx}] {random.choice(NOISE_BLOCKS_REALISTIC)}")
idx += 1
# Footer noise
if random.random() < 0.5:
for _ in range(random.randint(1, 2)):
blocks.append(f"[{idx}] {random.choice(NOISE_BLOCKS_REALISTIC)}")
idx += 1
if not relevant_indices or len(blocks) < 3:
continue
block_text = "\n".join(blocks)
# QE example
intervals = indices_to_intervals(relevant_indices)
query_type = row.get('query_type', 'general')
user_msg = f"URL: https://www.bing.com/search\nQuery: {query}\n\nBlocks:\n{block_text}\n\nOutput the index intervals of blocks relevant to the query."
all_examples.append({
"messages": [
{"role": "system", "content": SYSTEM_PROMPT_QE},
{"role": "user", "content": user_msg},
{"role": "assistant", "content": intervals}
],
"task_type": "query_relevant",
"source": f"msmarco_{query_type}"
})
except Exception as e:
continue
print(f" Total MS MARCO examples: {len(all_examples)}")
return all_examples
# ============================================================
# NA Examples (no relevant content)
# ============================================================
def create_na_examples(all_examples):
"""Create NA examples by mismatching queries with pages."""
print("Creating NA examples (mismatched query-page pairs)...")
# Get QE examples
qe_examples = [e for e in all_examples if e['task_type'] == 'query_relevant']
if len(qe_examples) < 100:
print(" Too few QE examples for NA generation")
return []
na_examples = []
for i in range(min(500, len(qe_examples) // 5)):
# Pick two random QE examples
idx_a = random.randint(0, len(qe_examples) - 1)
idx_b = (idx_a + random.randint(100, len(qe_examples) - 1)) % len(qe_examples)
# Use query from A, blocks from B
msgs_a = qe_examples[idx_a]['messages']
msgs_b = qe_examples[idx_b]['messages']
# Extract query from A
user_a = msgs_a[1]['content']
query_match = re.search(r'Query: (.+?)(\n|$)', user_a)
if not query_match:
continue
query = query_match.group(1).strip()
# Extract blocks from B
user_b = msgs_b[1]['content']
blocks_match = re.search(r'Blocks:\n(.+?)(\n\nOutput)', user_b, re.DOTALL)
if not blocks_match:
continue
blocks = blocks_match.group(1)
user_msg = f"URL: https://example.com\nQuery: {query}\n\nBlocks:\n{blocks}\n\nOutput the index intervals of blocks relevant to the query."
na_examples.append({
"messages": [
{"role": "system", "content": SYSTEM_PROMPT_QE},
{"role": "user", "content": user_msg},
{"role": "assistant", "content": "NA"}
],
"task_type": "query_relevant_na",
"source": "mismatched"
})
print(f" Created {len(na_examples)} NA examples")
return na_examples
# ============================================================
# Main Pipeline
# ============================================================
def main():
print("=" * 60)
print("Building DIVERSE IndexLM Training Data")
print("=" * 60)
# Collect from all sources
htmlrag_examples = load_htmlrag_data() # Real web HTML (primary)
multihoprag_examples = process_multihoprag() # News domain
hotpotqa_examples = process_hotpotqa() # Wikipedia (minority)
msmarco_examples = process_msmarco() # Diverse web QA
# Combine
all_examples = htmlrag_examples + multihoprag_examples + hotpotqa_examples + msmarco_examples
# Add NA examples
na_examples = create_na_examples(all_examples)
all_examples.extend(na_examples)
random.shuffle(all_examples)
# Print composition
print(f"\n{'='*60}")
print(f"Total examples: {len(all_examples)}")
source_counts = defaultdict(int)
type_counts = defaultdict(int)
for ex in all_examples:
source_counts[ex.get('source', 'unknown')] += 1
type_counts[ex['task_type']] += 1
print("\nBy source:")
for s, c in sorted(source_counts.items(), key=lambda x: -x[1]):
pct = c / len(all_examples) * 100
print(f" {s}: {c} ({pct:.1f}%)")
print("\nBy task type:")
for t, c in sorted(type_counts.items(), key=lambda x: -x[1]):
pct = c / len(all_examples) * 100
print(f" {t}: {c} ({pct:.1f}%)")
# Check token lengths
print("\nChecking token lengths...")
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
lengths = []
for ex in random.sample(all_examples, min(500, len(all_examples))):
text = tokenizer.apply_chat_template(ex['messages'], tokenize=False)
tokens = tokenizer.encode(text)
lengths.append(len(tokens))
print(f"Token length stats (sample of {len(lengths)}):")
print(f" Min: {min(lengths)}, Max: {max(lengths)}")
print(f" Mean: {sum(lengths)/len(lengths):.0f}, Median: {sorted(lengths)[len(lengths)//2]}")
# Filter by length
MAX_LEN = 4096
filtered = []
too_long = 0
for ex in all_examples:
text = tokenizer.apply_chat_template(ex['messages'], tokenize=False)
tokens = tokenizer.encode(text)
if len(tokens) <= MAX_LEN:
filtered.append(ex)
else:
too_long += 1
print(f"\nFiltered: {too_long} examples too long (>{MAX_LEN} tokens)")
print(f"Final dataset size: {len(filtered)}")
# Final composition
final_source_counts = defaultdict(int)
for ex in filtered:
final_source_counts[ex.get('source', 'unknown')] += 1
print("\nFinal composition by source:")
for s, c in sorted(final_source_counts.items(), key=lambda x: -x[1]):
pct = c / len(filtered) * 100
print(f" {s}: {c} ({pct:.1f}%)")
# Split
random.shuffle(filtered)
eval_size = min(500, len(filtered) // 10)
train_data = filtered[:-eval_size]
eval_data = filtered[-eval_size:]
print(f"\nTrain: {len(train_data)}, Eval: {len(eval_data)}")
# Create HF datasets
train_ds = Dataset.from_list([{"messages": ex["messages"]} for ex in train_data])
eval_ds = Dataset.from_list([{"messages": ex["messages"]} for ex in eval_data])
# Save locally
train_ds.save_to_disk("/app/indexlm_train_v2")
eval_ds.save_to_disk("/app/indexlm_eval_v2")
# Push to Hub
ds_dict = DatasetDict({"train": train_ds, "eval": eval_ds})
ds_dict.push_to_hub("OmAlve/indexlm-training-data", token=os.environ.get("HF_TOKEN"))
print(f"\n{'='*60}")
print("Done! Dataset pushed to OmAlve/indexlm-training-data")
print(f"{'='*60}")
if __name__ == "__main__":
main()

3
runs/Apr24_07-58-54_1eb67182ed08/events.out.tfevents.1777017534.1eb67182ed08.50747.0 Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:21a64bfa465482310059002822d2fb6e5ce46b4d17aaa650d28d737c73ab3c56
size 5623

3
runs/Apr24_08-03-09_1eb67182ed08/events.out.tfevents.1777017789.1eb67182ed08.51740.0 Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:f6a1e276745fd765b0a9bfafd89a29a6fb29e69b1dd62f21c69573f85769ffeb
size 6005

3
runs/Apr24_08-03-46_1eb67182ed08/events.out.tfevents.1777017826.1eb67182ed08.52486.0 Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:70c1d56045aa1b9bc45b1e326393dbb1f37b837654cdc29088ead54fa8106694
size 6008

3
runs/Apr24_08-06-45_1eb67182ed08/events.out.tfevents.1777018005.1eb67182ed08.53075.0 Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:5d9d339d3a50391043b225ee1dcbaec38d846f085ba065605cb9fef69a650d14
size 200909

3
tokenizer.json Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506
size 11422650

29
tokenizer_config.json Normal file
View File

@@ -0,0 +1,29 @@
{
"add_prefix_space": false,
"backend": "tokenizers",
"bos_token": null,
"clean_up_tokenization_spaces": false,
"eos_token": "<|im_end|>",
"errors": "replace",
"extra_special_tokens": [
"<|im_start|>",
"<|im_end|>",
"<|object_ref_start|>",
"<|object_ref_end|>",
"<|box_start|>",
"<|box_end|>",
"<|quad_start|>",
"<|quad_end|>",
"<|vision_start|>",
"<|vision_end|>",
"<|vision_pad|>",
"<|image_pad|>",
"<|video_pad|>"
],
"is_local": false,
"model_max_length": 131072,
"pad_token": "<|endoftext|>",
"split_special_tokens": false,
"tokenizer_class": "Qwen2Tokenizer",
"unk_token": null
}

137
train_indexlm.py Normal file
View File

@@ -0,0 +1,137 @@
"""
Reading Steiner - Fine-tune Qwen3-0.6B for Index-based Web Content Extraction
Based on: "An Index-based Approach for Efficient and Effective Web Content Extraction" (arxiv:2512.06641)
Base model: Qwen/Qwen3-0.6B (0.6B params, ideal for CPU deployment)
Training method: SFT with TRL SFTTrainer
Dataset: OmAlve/indexlm-training-data (21K+ multi-domain examples)
"""
import os
import torch
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import SFTTrainer, SFTConfig
import trackio
# ============ Configuration ============
MODEL_ID = "Qwen/Qwen3-0.6B"
DATASET_ID = "OmAlve/indexlm-training-data"
OUTPUT_DIR = "./reading-steiner"
HUB_MODEL_ID = "OmAlve/reading-steiner"
# Training hyperparameters (from paper: standard SFT)
LEARNING_RATE = 2e-5
NUM_EPOCHS = 3
BATCH_SIZE = 4
GRAD_ACCUM = 4 # Effective batch size = 16
MAX_SEQ_LENGTH = 4096
WARMUP_RATIO = 0.05
# ============ Setup Trackio ============
trackio.init(
name="reading-steiner-training",
project="reading-steiner"
)
# ============ Load Dataset ============
print("Loading dataset...")
dataset = load_dataset(DATASET_ID)
train_dataset = dataset["train"]
eval_dataset = dataset["eval"]
print(f"Train: {len(train_dataset)}, Eval: {len(eval_dataset)}")
# ============ Load Model & Tokenizer ============
print("Loading model and tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
# Ensure padding token is set
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(
MODEL_ID,
torch_dtype=torch.bfloat16,
attn_implementation="flash_attention_2", # Change to "sdpa" if flash-attn unavailable
)
print(f"Model loaded: {MODEL_ID}")
print(f"Model params: {sum(p.numel() for p in model.parameters()) / 1e6:.1f}M")
# ============ Training Config ============
training_args = SFTConfig(
output_dir=OUTPUT_DIR,
num_train_epochs=NUM_EPOCHS,
per_device_train_batch_size=BATCH_SIZE,
per_device_eval_batch_size=BATCH_SIZE,
gradient_accumulation_steps=GRAD_ACCUM,
learning_rate=LEARNING_RATE,
lr_scheduler_type="cosine",
warmup_ratio=WARMUP_RATIO,
weight_decay=0.01,
bf16=True,
gradient_checkpointing=True,
max_length=MAX_SEQ_LENGTH,
# Logging
logging_steps=10,
logging_first_step=True,
logging_strategy="steps",
disable_tqdm=True,
# Evaluation
eval_strategy="steps",
eval_steps=500,
# Saving
save_strategy="steps",
save_steps=500,
save_total_limit=3,
load_best_model_at_end=True,
metric_for_best_model="eval_loss",
greater_is_better=False,
# Hub push
push_to_hub=True,
hub_model_id=HUB_MODEL_ID,
hub_strategy="every_save",
# Performance
dataloader_num_workers=4,
dataloader_pin_memory=True,
# Report
report_to="none",
# Seed
seed=42,
)
# ============ Initialize Trainer ============
print("Initializing trainer...")
trainer = SFTTrainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
processing_class=tokenizer,
)
# ============ Train ============
print("Starting training...")
train_result = trainer.train()
# ============ Save Final Model ============
print("Saving final model...")
trainer.save_model(OUTPUT_DIR)
tokenizer.save_pretrained(OUTPUT_DIR)
# Push to Hub
print("Pushing to Hub...")
trainer.push_to_hub(commit_message="Final reading-steiner model")
# ============ Log Final Metrics ============
metrics = train_result.metrics
print(f"\nTraining complete!")
print(f" Train loss: {metrics.get('train_loss', 'N/A')}")
print(f" Train runtime: {metrics.get('train_runtime', 'N/A'):.0f}s")
print(f" Train samples/sec: {metrics.get('train_samples_per_second', 'N/A'):.1f}")
# Final eval
eval_metrics = trainer.evaluate()
print(f" Eval loss: {eval_metrics.get('eval_loss', 'N/A')}")
print(f"\nModel pushed to: https://huggingface.co/{HUB_MODEL_ID}")

3
training_args.bin Normal file
View File

@@ -0,0 +1,3 @@
version https://git-lfs.github.com/spec/v1
oid sha256:79a86123895b91f3924b594437b0ba747ae3d58fb7f8ab73d7c7d6f4be44b100
size 5304