thinhlpg
4 weeks ago
parent
504f0c6c8e
commit
333d1e596e
@ -0,0 +1,155 @@
|
|||||||
|
"""Prepares a deterministic sampled dev set (questions_dev.jsonl) from raw Musique dev data."""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import math
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
from collections import defaultdict
|
||||||
|
from pathlib import Path
|
||||||
|
|
||||||
|
|
||||||
|
def transform_musique_dev_data(input_path: str, output_path: str, sample_config: dict) -> None:
|
||||||
|
"""Transforms Musique dev data with deterministic stratified sampling using uniform selection from sorted lists.
|
||||||
|
|
||||||
|
Reads dev data, categorizes by hop type (2, 3, 4), sorts categories by ID,
|
||||||
|
selects N samples uniformly spaced from each sorted category based on sample_config,
|
||||||
|
combines, sorts final list by ID, combines answers/aliases, extracts supporting paras,
|
||||||
|
and writes the transformed data to output_path.
|
||||||
|
|
||||||
|
Args:
|
||||||
|
input_path: Path to the input JSONL file (e.g., data/raw/musique_ans_v1.0_dev.jsonl).
|
||||||
|
output_path: Path to the output JSONL file (e.g., data/processed/questions_dev.jsonl).
|
||||||
|
sample_config: Dictionary specifying samples per hop type (e.g., {"2hop": 20, "3hop": 15, "4hop": 15}).
|
||||||
|
"""
|
||||||
|
output_dir = Path(output_path).parent
|
||||||
|
output_dir.mkdir(parents=True, exist_ok=True)
|
||||||
|
|
||||||
|
print(f"Reading all data from {input_path} for dev sampling...")
|
||||||
|
all_data = []
|
||||||
|
try:
|
||||||
|
with open(input_path, "r", encoding="utf-8") as infile:
|
||||||
|
for line_num, line in enumerate(infile, 1):
|
||||||
|
try:
|
||||||
|
data = json.loads(line)
|
||||||
|
if "id" in data:
|
||||||
|
all_data.append(data)
|
||||||
|
else:
|
||||||
|
print(f"Warning: Skipping line {line_num} due to missing 'id' field in {input_path}")
|
||||||
|
except json.JSONDecodeError:
|
||||||
|
print(f"Warning: Skipping invalid JSON in line {line_num} of {input_path}")
|
||||||
|
except FileNotFoundError:
|
||||||
|
print(f"Error: Input file not found at {input_path}")
|
||||||
|
return
|
||||||
|
except Exception as e:
|
||||||
|
print(f"Error reading file {input_path}: {e}")
|
||||||
|
return
|
||||||
|
print(f"Read {len(all_data)} total samples from dev set.")
|
||||||
|
|
||||||
|
# Categorize data by hop count (2hop, 3hop, 4hop)
|
||||||
|
categorized_data = defaultdict(list)
|
||||||
|
print("Categorizing data by hop type (2, 3, 4)...")
|
||||||
|
for data in all_data:
|
||||||
|
q_id = data["id"]
|
||||||
|
hop_type = None
|
||||||
|
if q_id.startswith("2hop"):
|
||||||
|
hop_type = "2hop"
|
||||||
|
elif q_id.startswith("3hop"):
|
||||||
|
hop_type = "3hop"
|
||||||
|
elif q_id.startswith("4hop"):
|
||||||
|
hop_type = "4hop"
|
||||||
|
|
||||||
|
if hop_type:
|
||||||
|
categorized_data[hop_type].append(data)
|
||||||
|
|
||||||
|
# Deterministic sampling using sorting and uniform index selection
|
||||||
|
final_sample_list = []
|
||||||
|
total_target = sum(sample_config.values())
|
||||||
|
print(f"Sampling deterministically via uniform selection from sorted lists to get {total_target} dev samples...")
|
||||||
|
|
||||||
|
for hop_type, target_count in sample_config.items():
|
||||||
|
available_samples = categorized_data.get(hop_type, [])
|
||||||
|
current_count = len(available_samples)
|
||||||
|
print(f" {hop_type}: Found {current_count} samples, need {target_count}.")
|
||||||
|
|
||||||
|
if current_count == 0:
|
||||||
|
continue
|
||||||
|
|
||||||
|
available_samples.sort(key=lambda x: x["id"])
|
||||||
|
selected_samples_for_hop = []
|
||||||
|
if current_count < target_count:
|
||||||
|
print(f" Warning: Not enough samples for {hop_type}. Taking all {current_count} sorted samples.")
|
||||||
|
selected_samples_for_hop = available_samples
|
||||||
|
elif target_count > 0: # Ensure target_count is positive before selecting
|
||||||
|
print(f" Selecting {target_count} samples uniformly from {current_count}...")
|
||||||
|
# Calculate indices using integer interpretation of evenly spaced points
|
||||||
|
indices_to_take = [
|
||||||
|
int(i * (current_count - 1) / (target_count - 1)) if target_count > 1 else 0
|
||||||
|
for i in range(target_count)
|
||||||
|
] # Adjust index calc for edges
|
||||||
|
indices_to_take = sorted(list(set(indices_to_take))) # Ensure unique indices
|
||||||
|
# Simple fallback if uniqueness reduced count below target
|
||||||
|
while len(indices_to_take) < target_count and len(indices_to_take) < current_count:
|
||||||
|
next_val = indices_to_take[-1] + 1
|
||||||
|
if next_val < current_count:
|
||||||
|
indices_to_take.append(next_val)
|
||||||
|
else: # Cannot add more unique indices
|
||||||
|
break
|
||||||
|
selected_samples_for_hop = [
|
||||||
|
available_samples[idx] for idx in indices_to_take[:target_count]
|
||||||
|
] # Select based on unique indices, capped at target
|
||||||
|
|
||||||
|
final_sample_list.extend(selected_samples_for_hop)
|
||||||
|
|
||||||
|
print(f"Selected {len(final_sample_list)} dev samples in total.")
|
||||||
|
|
||||||
|
# Sort the final combined list by ID for consistent output order
|
||||||
|
print("Sorting the final combined dev sample list by ID...")
|
||||||
|
final_sample_list.sort(key=lambda x: x["id"])
|
||||||
|
|
||||||
|
# Process and write the selected samples
|
||||||
|
print(f"Processing and writing {len(final_sample_list)} selected dev samples to {output_path}...")
|
||||||
|
count = 0
|
||||||
|
try:
|
||||||
|
with open(output_path, "w", encoding="utf-8") as outfile:
|
||||||
|
for data in final_sample_list:
|
||||||
|
try:
|
||||||
|
supporting_paragraphs = [
|
||||||
|
p["paragraph_text"] for p in data.get("paragraphs", []) if p.get("is_supporting", False)
|
||||||
|
]
|
||||||
|
main_answer = data.get("answer", "")
|
||||||
|
aliases = data.get("answer_aliases", [])
|
||||||
|
all_answers = [main_answer] + (aliases if isinstance(aliases, list) else [])
|
||||||
|
valid_answers = [str(ans).strip() for ans in all_answers if ans and str(ans).strip()]
|
||||||
|
unique_valid_answers = list(set(valid_answers)) # Keep unique, don't sort alphabetically
|
||||||
|
combined_answer_str = " OR ".join(unique_valid_answers)
|
||||||
|
|
||||||
|
output_data = {
|
||||||
|
"id": data.get("id"),
|
||||||
|
"question": data.get("question"),
|
||||||
|
"answer": combined_answer_str,
|
||||||
|
"supporting_paragraphs": supporting_paragraphs,
|
||||||
|
}
|
||||||
|
outfile.write(json.dumps(output_data) + "\n")
|
||||||
|
count += 1
|
||||||
|
except KeyError as e:
|
||||||
|
print(f"Skipping sample due to missing key {e}: {data.get('id')}")
|
||||||
|
print(f"Successfully processed and wrote {count} dev samples.")
|
||||||
|
except Exception as e:
|
||||||
|
print(f"An unexpected error occurred during writing: {e}")
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == "__main__":
|
||||||
|
# Define file paths relative to the project root
|
||||||
|
# Ensure this script is run from the project root or adjust paths accordingly
|
||||||
|
RAW_DIR = Path("data/raw")
|
||||||
|
PROCESSED_DIR = Path("data/processed")
|
||||||
|
|
||||||
|
# Define sampling configuration for the dev set
|
||||||
|
DEV_SAMPLING_CONFIG = {"2hop": 20, "3hop": 15, "4hop": 15} # Total = 50
|
||||||
|
|
||||||
|
INPUT_FILE = RAW_DIR / "musique_ans_v1.0_dev.jsonl"
|
||||||
|
OUTPUT_FILE = PROCESSED_DIR / "questions_dev.jsonl"
|
||||||
|
|
||||||
|
transform_musique_dev_data(str(INPUT_FILE), str(OUTPUT_FILE), DEV_SAMPLING_CONFIG)
|
||||||
|
|
||||||
|
print(f"\nMusique DEV JSONL transformation and deterministic sampling complete.")
|
||||||
Loading…
Reference in new issue