main.py

import json
import os
import numpy as np
import random

from data.dataset import load_data, make_vocabulary, compute_keep_prob, build_noise_dist
from eval import run_benchmarks
import model.cbow as cbow_model
import model.sgns as sgns_model
from utils.args import parse_args
from utils.checkpoint import save_final


if __name__ == "__main__":

    # Parse arguments
    args = parse_args()

    # Set random seed for reproducibility
    random.seed(args.seed)
    np.random.seed(args.seed)

    # Load dataset and build vocabulary
    words = load_data()
    data, idx_to_word, word_to_idx, vocab_words, freq = make_vocabulary(words, args)

    # Compute subsampling probabilities and noise distribution for negative sampling
    keep_prob, word_freqs = compute_keep_prob(freq, idx_to_word, len(vocab_words), args)
    noise_dist = build_noise_dist(word_freqs)

    # Select model module
    m = sgns_model if args.model == "sgns" else cbow_model
    print(f"Model: {args.model.upper()}")

    # Initialize model embeddings
    W_in, W_out = m.build_model(len(vocab_words), args.embed_dim, args.seed)

    # Set up mid-training eval callback if enabled
    eval_callback = None
    if args.eval_interval > 0 and (args.benchmark_analogy or args.benchmark_similarity):
        out_dir = os.path.dirname(args.w_in_path) or "."
        eval_log_path = os.path.join(out_dir, "eval_log.json")
        eval_records = []

        def eval_callback(epoch, W_in):
            print(f"\n--- Mid-training eval (epoch {epoch}) ---")
            result = run_benchmarks(
                W_in, word_to_idx, idx_to_word,
                args.benchmark_analogy, args.benchmark_similarity,
            )
            result["epoch"] = epoch
            eval_records.append(result)
            os.makedirs(out_dir, exist_ok=True)
            with open(eval_log_path, "w") as f:
                json.dump(eval_records, f, indent=2)

    # Train the model
    m.train(data, keep_prob, noise_dist, W_in, W_out, args, eval_callback=eval_callback)

    # Save final embeddings and vocabulary to disk
    save_final(W_in, W_out, idx_to_word, vocab_words, args)