folksy_idioms/scripts/format_training_pairs.py

385 lines
14 KiB
Python

#!/usr/bin/env python3
"""Format filtered sayings into training pairs for fine-tuning.
Each polished saying generates 3-5 training pairs with different input framings.
Also generates fictional entity training pairs.
Usage:
python scripts/format_training_pairs.py
python scripts/format_training_pairs.py --input corpus/corpus_filtered.jsonl --output corpus/training_pairs.jsonl
"""
import argparse
import csv
import json
import random
import sys
from pathlib import Path
SCRIPT_DIR = Path(__file__).parent
PROJECT_DIR = SCRIPT_DIR.parent
CORPUS_DIR = PROJECT_DIR / "corpus"
DATA_DIR = PROJECT_DIR / "data"
EXAMPLES_DIR = PROJECT_DIR / "examples"
# Template name mappings for human-readable prompts
TEMPLATE_NAMES = {
"deconstruction": "deconstruction",
"denial_of_consequences": "denial of consequences",
"ironic_deficiency": "ironic deficiency",
"futile_preparation": "futile preparation",
"hypocritical_complaint": "hypocritical complaint",
"tautological_wisdom": "tautological wisdom",
"false_equivalence": "false equivalence",
}
PERSONAS = ["farmer", "grandmother", "old sailor", "blacksmith", "innkeeper", "shepherd"]
OPEN_ENDED_PROMPTS = [
"Tell me some folk wisdom.",
"What do they say?",
"Give me a proverb.",
"Share some old-time wisdom.",
"What's a good saying?",
]
# Auto-generated fictional entities for additional training pairs
AUTO_ENTITIES = [
{
"name": "Stoneclaw",
"categories": ["animal", "predator"],
"properties": ["fierce", "rocky", "nocturnal"],
"relations": {"AtLocation": ["cave", "mountain"], "HasA": ["claws", "scales"], "CapableOf": ["hunting", "climbing"]},
},
{
"name": "Duskmelon",
"categories": ["fruit", "food"],
"properties": ["purple", "sweet", "fragrant"],
"relations": {"AtLocation": ["garden", "market"], "UsedFor": ["eating", "jam"], "MadeOf": ["seed", "juice"]},
},
{
"name": "Windloom",
"categories": ["tool", "craft"],
"properties": ["wooden", "portable", "intricate"],
"relations": {"UsedFor": ["weaving", "thread"], "MadeOf": ["wood", "string"], "AtLocation": ["workshop", "cottage"]},
},
{
"name": "Briarvine",
"categories": ["plant", "herb"],
"properties": ["thorny", "green", "medicinal"],
"relations": {"AtLocation": ["forest", "hedge"], "UsedFor": ["healing", "tea"], "HasA": ["thorn", "leaf"]},
},
{
"name": "Mudhog",
"categories": ["animal", "livestock"],
"properties": ["muddy", "stubborn", "heavy"],
"relations": {"AtLocation": ["farm", "swamp"], "Desires": ["food", "mud"], "CapableOf": ["digging", "rooting"]},
},
{
"name": "Frostberry",
"categories": ["fruit", "food"],
"properties": ["cold", "blue", "tiny"],
"relations": {"AtLocation": ["mountain", "tundra"], "UsedFor": ["eating", "preserves"], "HasProperty": ["cold", "tart"]},
},
{
"name": "Lanternmoss",
"categories": ["plant", "fungus"],
"properties": ["glowing", "damp", "soft"],
"relations": {"AtLocation": ["cave", "swamp"], "UsedFor": ["light", "decoration"], "HasProperty": ["luminous", "fragile"]},
},
{
"name": "Cinderhawk",
"categories": ["bird", "animal"],
"properties": ["fiery", "fast", "red"],
"relations": {"AtLocation": ["mountain", "volcano"], "CapableOf": ["flying", "hunting"], "HasA": ["talons", "feathers"]},
},
{
"name": "Rootstone",
"categories": ["stone", "material"],
"properties": ["veined", "hard", "ancient"],
"relations": {"AtLocation": ["quarry", "riverbed"], "UsedFor": ["building", "carving"], "MadeOf": ["mineral", "root"]},
},
{
"name": "Silkwort",
"categories": ["plant", "fiber"],
"properties": ["silky", "white", "tall"],
"relations": {"AtLocation": ["field", "meadow"], "UsedFor": ["weaving", "cloth"], "HasA": ["stem", "fiber"]},
},
{
"name": "Kettlefrog",
"categories": ["animal", "amphibian"],
"properties": ["loud", "round", "green"],
"relations": {"AtLocation": ["pond", "marsh"], "CapableOf": ["jumping", "croaking"], "Desires": ["flies", "water"]},
},
{
"name": "Dustwheat",
"categories": ["crop", "grain"],
"properties": ["dry", "golden", "hardy"],
"relations": {"AtLocation": ["field", "barn"], "UsedFor": ["bread", "flour"], "HasPrerequisite": ["rain", "soil"]},
},
]
def format_entity_description(entity):
"""Format entity into a natural description string."""
name = entity["name"]
cats = entity.get("categories", [])
props = entity.get("properties", [])
rels = entity.get("relations", {})
parts = []
# Category description
if props and cats:
prop_str = ", ".join(props[:3])
cat_str = " and ".join(cats[:2])
parts.append(f"A {name} is a {prop_str} {cat_str}.")
elif cats:
parts.append(f"A {name} is a {' and '.join(cats[:2])}.")
# Location
if "AtLocation" in rels:
locs = rels["AtLocation"]
parts.append(f"It is found near {' and '.join(locs[:2])}.")
# Parts/properties
if "HasA" in rels:
has = rels["HasA"]
parts.append(f"It has {', '.join(has[:3])}.")
# Capabilities
if "CapableOf" in rels:
caps = rels["CapableOf"]
parts.append(f"It can {' and '.join(caps[:2])}.")
# Uses
if "UsedFor" in rels:
uses = rels["UsedFor"]
parts.append(f"It is used for {' and '.join(uses[:2])}.")
return " ".join(parts)
def load_vocab_categories():
"""Load vocab to get word -> categories mapping."""
word_cats = {}
vocab_path = DATA_DIR / "folksy_vocab.csv"
if vocab_path.exists():
with open(vocab_path, newline="", encoding="utf-8") as f:
for row in csv.DictReader(f):
word = row["word"]
cats = [c.strip() for c in row["categories"].split(",") if c.strip()]
word_cats[word] = cats
return word_cats
def generate_training_pairs(entry, word_cats):
"""Generate 3-5 training pairs for a single polished saying."""
polished = entry.get("polished_text", "")
slots = entry.get("slots", {})
meta_template = entry.get("meta_template", "")
# Collect source words (concrete nouns from slots)
source_words = [v for v in slots.values()
if v and not v.startswith("a ") and not v.startswith("an ") and len(v) > 1]
# Determine categories of slot words
slot_categories = set()
for word in source_words:
word_lower = word.lower().replace(" ", "_")
if word_lower in word_cats:
slot_categories.update(word_cats[word_lower])
pairs = []
base = {
"output": polished,
"meta_template": meta_template,
"source_words": source_words,
}
# 1. Word-seeded (always include)
if source_words:
word = random.choice(source_words)
pairs.append({**base, "input": f"Tell me something about {word}."})
# 2. Category-seeded (always include if we have categories)
if slot_categories:
cat = random.choice(list(slot_categories))
pairs.append({**base, "input": f"Tell me a saying about {cat}."})
# 3. Persona-seeded (always include)
persona = random.choice(PERSONAS)
if source_words:
word = random.choice(source_words)
pairs.append({**base, "input": f"What would a {persona} say about {word}?"})
# 4. Template-seeded (include ~70% of the time)
if random.random() < 0.7:
template_name = TEMPLATE_NAMES.get(meta_template, meta_template)
pairs.append({**base, "input": f"Give me a {template_name} proverb."})
# 5. Open-ended (include ~30% of the time)
if random.random() < 0.3:
prompt = random.choice(OPEN_ENDED_PROMPTS)
pairs.append({**base, "input": prompt})
return pairs
def generate_fictional_pairs(entities):
"""Generate training pairs for fictional entities.
These pairs include the entity description in the input.
"""
pairs = []
# Generate 15-25 pairs per entity
for entity in entities:
name = entity["name"]
desc = format_entity_description(entity)
props = entity.get("properties", [])
rels = entity.get("relations", {})
# Collect words related to this entity
related_words = []
for targets in rels.values():
related_words.extend(targets)
n_pairs = random.randint(15, 25)
for _ in range(n_pairs):
framing = random.choice(["persona", "word", "category", "open"])
if framing == "persona":
persona = random.choice(PERSONAS)
input_text = f"{desc} What would a {persona} say about a {name}?"
elif framing == "word" and related_words:
word = random.choice(related_words)
input_text = f"{desc} Tell me a saying about {name} and {word}."
elif framing == "category":
cats = entity.get("categories", ["thing"])
cat = random.choice(cats)
input_text = f"{desc} Give me folk wisdom about this {cat}."
else:
input_text = f"{desc} Tell me some folk wisdom about {name}."
# Placeholder output — these would ideally be generated through the
# template engine with fictional entities loaded, then polished.
# For now, generate a structural placeholder that indicates the
# entity relationships.
pairs.append({
"input": input_text,
"output": "", # Will be filled by actual generation
"meta_template": "fictional",
"source_words": [name] + related_words[:3],
"_needs_generation": True,
"_entity": entity,
})
return pairs
def main():
parser = argparse.ArgumentParser(description="Format training pairs for fine-tuning.")
parser.add_argument("--input", default=str(CORPUS_DIR / "corpus_filtered.jsonl"),
help="Input filtered JSONL file")
parser.add_argument("--output", default=str(CORPUS_DIR / "training_pairs.jsonl"),
help="Output training pairs JSONL file")
parser.add_argument("--entities", default=str(EXAMPLES_DIR / "my_world.json"),
help="Fictional entities JSON file")
args = parser.parse_args()
input_path = Path(args.input)
output_path = Path(args.output)
entities_path = Path(args.entities)
if not input_path.exists():
print(f"Error: {input_path} not found.", file=sys.stderr)
sys.exit(1)
# Load vocab categories
word_cats = load_vocab_categories()
# Load filtered entries
entries = []
with open(input_path, encoding="utf-8") as f:
for line in f:
line = line.strip()
if line:
entries.append(json.loads(line))
print(f"Loaded {len(entries)} filtered entries")
# Generate training pairs for each entry
all_pairs = []
for entry in entries:
pairs = generate_training_pairs(entry, word_cats)
all_pairs.extend(pairs)
print(f"Generated {len(all_pairs)} training pairs from polished sayings")
# Generate fictional entity pairs
fictional_entities = []
if entities_path.exists():
with open(entities_path, encoding="utf-8") as f:
data = json.load(f)
fictional_entities = data.get("entities", [])
print(f"Loaded {len(fictional_entities)} fictional entities from {entities_path}")
# Add auto-generated entities
fictional_entities.extend(AUTO_ENTITIES)
print(f"Total fictional entities (file + auto-generated): {len(fictional_entities)}")
fictional_pairs = generate_fictional_pairs(fictional_entities)
# Filter out placeholder pairs (those that still need generation)
# In a full pipeline, these would be generated through the template engine.
# For now, skip any with empty output.
real_fictional = [p for p in fictional_pairs if p.get("output")]
placeholder_fictional = [p for p in fictional_pairs if not p.get("output")]
if placeholder_fictional:
print(f" {len(placeholder_fictional)} fictional pairs need generation via template engine")
print(f" (Run folksy_generator.py with --entities to generate these, then re-run this script)")
all_pairs.extend(real_fictional)
# Clean up internal fields before writing
for pair in all_pairs:
pair.pop("_needs_generation", None)
pair.pop("_entity", None)
# Write output
output_path.parent.mkdir(parents=True, exist_ok=True)
with open(output_path, "w", encoding="utf-8") as f:
for pair in all_pairs:
f.write(json.dumps(pair, ensure_ascii=False) + "\n")
# Stats
from collections import Counter
input_types = Counter()
for pair in all_pairs:
inp = pair["input"]
if inp.startswith("Tell me something about"):
input_types["word_seeded"] += 1
elif inp.startswith("Tell me a saying about"):
input_types["category_seeded"] += 1
elif inp.startswith("What would a"):
input_types["persona_seeded"] += 1
elif inp.startswith("Give me a") and "proverb" in inp:
input_types["template_seeded"] += 1
elif any(inp.startswith(p) for p in ["Tell me some folk", "What do they", "Give me a proverb", "Share some", "What's a good"]):
input_types["open_ended"] += 1
else:
input_types["fictional"] += 1
print(f"\nTotal training pairs: {len(all_pairs)}")
print("Distribution by input type:")
for itype, count in sorted(input_types.items()):
print(f" {itype:20s} {count:5d}")
print(f"\nOutput: {output_path}")
if __name__ == "__main__":
main()