train.py

from datasets import load_dataset, Dataset
from gcsfs import GCSFileSystem
from transformers import (
    AutoTokenizer,
    GPT2Config,
    GPT2LMHeadModel,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling
)
import os
import json
from google.colab import auth
import google.auth

# --- GCS Setup ---
auth.authenticate_user()

# Initialize GCSFS with credentials
credentials, project = google.auth.default()
fs = GCSFileSystem(project=project, token=credentials)

# Define the correct path based on your screenshot

bucket_name = "lean-universe"


# List all JSONL files in that path
file_paths = fs.glob(f"{bucket_name}/proof_model_data.jsonl")
print(f"Found {len(file_paths)} files: {file_paths}")
print(f"Found {len(file_paths)} files: {file_paths}")

# Check if we found any files
if not file_paths:
    # Try without the subdirectory
    prefix = "huggingface-datasets/"
    file_paths = fs.glob(f"{bucket_name}/{prefix}/*.jsonl")
    print(f"Retrying with broader path, found {len(file_paths)} files")

    if not file_paths:
        # Try the root of the bucket
        file_paths = fs.glob(f"{bucket_name}/*.jsonl")
        print(f"Searching root of bucket, found {len(file_paths)} files")

# If still no files, try to list all files in the bucket to see what's available
if not file_paths:
    all_files = fs.ls(bucket_name)
    print(f"All files in bucket: {all_files}")

    # Look specifically for train-chunk files seen in the screenshot
    train_chunk_files = [f for f in all_files if "train-chunk" in f]
    if train_chunk_files:
        file_paths = train_chunk_files
        print(f"Found train-chunk files: {file_paths}")

# Load and combine the JSONL files
all_data = []
for file_path in file_paths:
    try:
        with fs.open(file_path, 'r') as f:
            file_data = []
            for i, line in enumerate(f):
                try:
                    file_data.append(json.loads(line))
                except json.JSONDecodeError:
                    print(f"Error decoding JSON in {file_path} at line {i+1}")

            print(f"Loaded {len(file_data)} records from {file_path}")

            # If we have data, check the first item's structure
            if file_data:
                print(f"First item keys: {list(file_data[0].keys())}")
                all_data.extend(file_data)
    except Exception as e:
        print(f"Error loading file {file_path}: {e}")

print(f"Total records loaded: {len(all_data)}")

# If we have no data, create a small synthetic dataset for testing
if not all_data:
    print("No data found. Creating synthetic dataset for testing...")
    all_data = [
        {
            "natural_language_statement": "The sum of two even numbers is even",
            "formal_statement": "forall a b, even a -> even b -> even (a + b)",
            "tactic": "intros a b Ha Hb. unfold even in *. destruct Ha as [k Hk]. destruct Hb as [j Hj]. exists (k + j). rewrite Hk, Hj. ring."
        },
        {
            "natural_language_statement": "The product of any number with zero is zero",
            "formal_statement": "forall a, a * 0 = 0",
            "tactic": "intro a. ring."
        }
    ]
    print("Created synthetic dataset with 2 examples")

# Create a Hugging Face dataset from the loaded data
# First check what fields are available
if all_data:
    available_keys = set()
    for item in all_data:
        available_keys.update(item.keys())
    print(f"Available fields in the data: {available_keys}")

# Determine which fields to use
field_mapping = {
    "natural_language_statement": ["natural_language_statement", "nl_statement", "statement"],
    "formal_statement": ["formal_statement", "theorem", "formal"],
    "tactic": ["tactic", "proof", "tactics"]
}

dataset_dict = {}
for target_field, possible_fields in field_mapping.items():
    for field in possible_fields:
        if any(field in item for item in all_data):
            dataset_dict[target_field] = [item.get(field, "") for item in all_data]
            print(f"Using '{field}' for '{target_field}'")
            break
    else:
        # If no matching field found, use empty strings
        dataset_dict[target_field] = ["" for _ in all_data]
        print(f"No field found for '{target_field}', using empty strings")

# Create the dataset
dataset = Dataset.from_dict(dataset_dict)
print(f"Created dataset with {len(dataset)} examples")

# Print a sample to verify
print("\nSample from the dataset:")
if len(dataset) > 0:
    print(dataset[0])

# Create a train/test split
dataset_dict = dataset.train_test_split(test_size=0.1, seed=42)
print(f"Train set: {len(dataset_dict['train'])} examples")
print(f"Test set: {len(dataset_dict['test'])} examples")

# Initialize tokenizer
tokenizer = AutoTokenizer.from_pretrained("gpt2", use_fast=True)
special_tokens = {
    "additional_special_tokens": [
        "<theorem>", "</theorem>",
        "<proof>", "</proof>",
        "<tactic>", "</tactic>",
        "<step>", "</step>",
        "<natural>", "</natural>",
        "<formal>", "</formal>"
    ],
    "pad_token": "[PAD]"
}
tokenizer.add_special_tokens(special_tokens)

# Configure model
config = GPT2Config(
    vocab_size=len(tokenizer),
    n_positions=1024,  # Reduced context length for testing
    n_embd=384,
    n_layer=6,  # Reduced for testing
    n_head=6,  # Reduced for testing
    bos_token_id=tokenizer.bos_token_id,
    eos_token_id=tokenizer.eos_token_id,
    pad_token_id=tokenizer.pad_token_id,
    resid_pdrop=0.1
)
model = GPT2LMHeadModel(config)
model.resize_token_embeddings(len(tokenizer))

# Preprocessing
def format_examples(examples):
    formatted_texts = []
    for nl_stmt, formal_stmt, tactic in zip(
        examples["natural_language_statement"],
        examples["formal_statement"],
        examples["tactic"]
    ):
        formatted = (
            f"<natural>{nl_stmt}</natural>\n"
            f"<formal>{formal_stmt}</formal>\n"
            f"<proof>{tactic}</proof>"
        )
        formatted_texts.append(formatted)
    return {"text": formatted_texts}

def tokenize_function(examples):
    outputs = tokenizer(
        examples["text"],
        truncation=True,
        max_length=512,  # Reduced for testing
        padding="max_length",
        return_tensors=None  # Changed from "pt" to None
    )
    return outputs

# Process dataset
print("Processing train dataset...")
dataset_dict["train"] = dataset_dict["train"].map(format_examples, batched=True)
print("Processing test dataset...")
dataset_dict["test"] = dataset_dict["test"].map(format_examples, batched=True)

print("Tokenizing train dataset...")
tokenized_train = dataset_dict["train"].map(
    tokenize_function,
    batched=True,
    remove_columns=dataset_dict["train"].column_names
)

print("Tokenizing test dataset...")
tokenized_test = dataset_dict["test"].map(
    tokenize_function,
    batched=True,
    remove_columns=dataset_dict["test"].column_names
)

print(f"Tokenized train dataset has {len(tokenized_train)} examples")
print(f"Tokenized test dataset has {len(tokenized_test)} examples")

# Print column names to verify
print(f"Columns in tokenized train dataset: {tokenized_train.column_names}")

# Check if dataset is empty
if len(tokenized_train) == 0:
    raise ValueError("Training dataset is empty!")

# Training setup
training_args = TrainingArguments(
    output_dir="./lean4-workbook-slm",
    num_train_epochs=2,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    gradient_accumulation_steps=2,
    learning_rate=5e-5,
    weight_decay=0.01,
    warmup_ratio=0.1,
    evaluation_strategy="steps",
    eval_steps=100,
    save_steps=500,
    logging_steps=10,
    load_best_model_at_end=True,
    metric_for_best_model="eval_loss",
    greater_is_better=False,
    # Disable wandb reporting
    report_to="none"  # Change this line from "tensorboard" to "none"
)

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=False
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_train,
    eval_dataset=tokenized_test,
    data_collator=data_collator
)

# Start training
print("Starting training...")
trainer.train()

# Save model
model.save_pretrained("./lean4-theorem-prover")
tokenizer.save_pretrained("./lean4-theorem-prover")