898 lines
35 KiB
Python
898 lines
35 KiB
Python
"""
|
|
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()
|