BiGRU_T_version / scripts /train_fast.py

Upload folder using huggingface_hub

d855e28 verified 3 months ago

5.02 kB

#!/usr/bin/env python3
"""train_fast.py — Treino rápido de bug-detection (2 épocas, modelo pequeno).
Versão otimizada para rodar em < 5 minutos no ambiente Xeon 2-core:
- max_modules=4 (em vez de 8)
- bigru_hidden=16 (em vez de 32)
- d_transformer=32 (em vez de 64)
- num_layers_train=1, num_layers_hyp=1
- max_samples=8 por dataset
- 1 dataset (CEIA-POSITIVO/ultrachat_br_clustred_balanced_v1)
Valida todos os 4 lemas sem pretender produzir modelo SOTA.
"""
from __future__ import annotations

import os
import sys
import time
import json
import logging
from pathlib import Path

# Adiciona src/ ao path
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

# Otimização Xeon
os.environ.setdefault("OMP_NUM_THREADS", "2")
os.environ.setdefault("MKL_NUM_THREADS", "2")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import torch
torch.set_num_threads(2)

try:
    from bigru_t.utils.xeon_runtime import optimize_xeon_environment
    optimize_xeon_environment()
except Exception as e:
    logging.warning(f"Could not apply Xeon optimization: {e}")

from bigru_t import (
    UnifiedModel, UnifiedModelConfig, create_unified_model,
    BiGRU_T_Trainer, TrainerConfig,
)
from bigru_t.data.streaming_datasets import stream_dataset

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
    handlers=[logging.StreamHandler()],
)
logger = logging.getLogger(__name__)


def main():
    # HF token
    hf_token = os.environ.get("HF_TOKEN") or None

    # Modelo pequeno para bug-detection rápido
    model_config = UnifiedModelConfig(
        vocab_size=16384,
        d_model=64,
        max_seq_len=32,
        pad_token_id=1,
        max_modules=4,           # 4 módulos em vez de 8
        bigru_hidden=16,         # 16 em vez de 32
        d_transformer=32,        # 32 em vez de 64
        nhead_tu=2,              # 2 em vez de 4 (32/2=16 ok)
        d_ff_tu=64,              # 64 em vez de 128
        output_dim_u8cell=32,    # 32 em vez de 64
        cache_len=8,             # 8 em vez de 16
        d_cache=64,              # 64 em vez de 128
        nhead_orq=2,             # 2 em vez de 8 (64/2=32 ok)
        d_ff_orq=128,            # 128 em vez de 256
        trainT_dim=64,
        nhead_train=2,
        d_ff_train=128,
        num_layers_train=1,      # 1 em vez de 2
        hypT_dim=64,
        nhead_hyp=2,
        d_ff_hyp=128,
        num_layers_hyp=1,        # 1 em vez de 2
        num_bits=8,
        dropout=0.1,
    )

    logger.info("Criando UnifiedModel (small config)...")
    model, _ = create_unified_model(model_config)
    params = model.count_parameters()
    logger.info(f"  params: {params['total']:,} ({params['total_M']:.2f}M)")

    # Tokenizer
    from tokenizers import Tokenizer
    tok_path = "/home/z/my-project/source/model_final/tokenizer/tokenizer.json"
    tokenizer = Tokenizer.from_file(tok_path)

    # Carrega 1 dataset
    ds_name = "CEIA-POSITIVO/ultrachat_br_clustred_balanced_v1"
    max_samples = 8
    logger.info(f"Loading {max_samples} samples from {ds_name}...")
    samples = []
    t0 = time.time()
    for sample in stream_dataset(ds_name, max_samples=max_samples, hf_token=hf_token):
        samples.append(sample)
        if len(samples) >= max_samples:
            break
        if time.time() - t0 > 120:
            logger.warning(f"Timeout após {len(samples)} amostras")
            break
    logger.info(f"Loaded {len(samples)} samples in {time.time()-t0:.1f}s")

    if not samples:
        logger.error("Nenhuma amostra carregada")
        sys.exit(1)

    # Split 90/10
    import random
    random.seed(42)
    random.shuffle(samples)
    split = max(1, int(0.9 * len(samples)))
    train_samples = samples[:split]
    val_samples = samples[split:] or samples[:1]
    logger.info(f"  train: {len(train_samples)} | val: {len(val_samples)}")

    # Trainer config
    trainer_config = TrainerConfig(
        epochs=2,
        datasets=ds_name,
        max_samples_per_dataset=max_samples,
        max_seq_len=32,
        per_device_batch_size=1,
        grad_accum=2,
        lr=1e-3,
        use_hypothesis=True,
        meta_interval=5,
        log_every=1,
        save_temp_every=100,  # não salvar temp (modelo pequeno)
        output_dir="/home/z/my-project/BiGRU_T_version/model_final",
        temp_dir="/home/z/my-project/BiGRU_T_version/_temp_checkpoints",
        keep_temp=False,
    )

    trainer = BiGRU_T_Trainer(
        model=model,
        tokenizer=tokenizer,
        config=trainer_config,
        train_samples=train_samples,
        val_samples=val_samples,
    )

    result = trainer.train()

    print("\n=== Resultado Final ===")
    print(json.dumps(result, indent=2, default=str))

    # Apaga HF token
    if "HF_TOKEN" in os.environ:
        del os.environ["HF_TOKEN"]
        logger.info("HF_TOKEN apagado do ambiente")

    sys.exit(0 if not result.get("killed") else 1)


if __name__ == "__main__":
    main()

	#!/usr/bin/env python3
	"""train_fast.py — Treino rápido de bug-detection (2 épocas, modelo pequeno).

	Versão otimizada para rodar em < 5 minutos no ambiente Xeon 2-core:
	- max_modules=4 (em vez de 8)
	- bigru_hidden=16 (em vez de 32)
	- d_transformer=32 (em vez de 64)
	- num_layers_train=1, num_layers_hyp=1
	- max_samples=8 por dataset
	- 1 dataset (CEIA-POSITIVO/ultrachat_br_clustred_balanced_v1)

	Valida todos os 4 lemas sem pretender produzir modelo SOTA.
	"""
	from __future__ import annotations

	import os
	import sys
	import time
	import json
	import logging
	from pathlib import Path

	# Adiciona src/ ao path
	sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

	# Otimização Xeon
	os.environ.setdefault("OMP_NUM_THREADS", "2")
	os.environ.setdefault("MKL_NUM_THREADS", "2")
	os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

	import torch
	torch.set_num_threads(2)

	try:
	from bigru_t.utils.xeon_runtime import optimize_xeon_environment
	optimize_xeon_environment()
	except Exception as e:
	logging.warning(f"Could not apply Xeon optimization: {e}")

	from bigru_t import (
	UnifiedModel, UnifiedModelConfig, create_unified_model,
	BiGRU_T_Trainer, TrainerConfig,
	)
	from bigru_t.data.streaming_datasets import stream_dataset

	logging.basicConfig(
	level=logging.INFO,
	format="%(asctime)s [%(levelname)s] %(message)s",
	datefmt="%H:%M:%S",
	handlers=[logging.StreamHandler()],
	)
	logger = logging.getLogger(__name__)


	def main():
	# HF token
	hf_token = os.environ.get("HF_TOKEN") or None

	# Modelo pequeno para bug-detection rápido
	model_config = UnifiedModelConfig(
	vocab_size=16384,
	d_model=64,
	max_seq_len=32,
	pad_token_id=1,
	max_modules=4, # 4 módulos em vez de 8
	bigru_hidden=16, # 16 em vez de 32
	d_transformer=32, # 32 em vez de 64
	nhead_tu=2, # 2 em vez de 4 (32/2=16 ok)
	d_ff_tu=64, # 64 em vez de 128
	output_dim_u8cell=32, # 32 em vez de 64
	cache_len=8, # 8 em vez de 16
	d_cache=64, # 64 em vez de 128
	nhead_orq=2, # 2 em vez de 8 (64/2=32 ok)
	d_ff_orq=128, # 128 em vez de 256
	trainT_dim=64,
	nhead_train=2,
	d_ff_train=128,
	num_layers_train=1, # 1 em vez de 2
	hypT_dim=64,
	nhead_hyp=2,
	d_ff_hyp=128,
	num_layers_hyp=1, # 1 em vez de 2
	num_bits=8,
	dropout=0.1,
	)

	logger.info("Criando UnifiedModel (small config)...")
	model, _ = create_unified_model(model_config)
	params = model.count_parameters()
	logger.info(f" params: {params['total']:,} ({params['total_M']:.2f}M)")

	# Tokenizer
	from tokenizers import Tokenizer
	tok_path = "/home/z/my-project/source/model_final/tokenizer/tokenizer.json"
	tokenizer = Tokenizer.from_file(tok_path)

	# Carrega 1 dataset
	ds_name = "CEIA-POSITIVO/ultrachat_br_clustred_balanced_v1"
	max_samples = 8
	logger.info(f"Loading {max_samples} samples from {ds_name}...")
	samples = []
	t0 = time.time()
	for sample in stream_dataset(ds_name, max_samples=max_samples, hf_token=hf_token):
	samples.append(sample)
	if len(samples) >= max_samples:
	break
	if time.time() - t0 > 120:
	logger.warning(f"Timeout após {len(samples)} amostras")
	break
	logger.info(f"Loaded {len(samples)} samples in {time.time()-t0:.1f}s")

	if not samples:
	logger.error("Nenhuma amostra carregada")
	sys.exit(1)

	# Split 90/10
	import random
	random.seed(42)
	random.shuffle(samples)
	split = max(1, int(0.9 * len(samples)))
	train_samples = samples[:split]
	val_samples = samples[split:] or samples[:1]
	logger.info(f" train: {len(train_samples)} \| val: {len(val_samples)}")

	# Trainer config
	trainer_config = TrainerConfig(
	epochs=2,
	datasets=ds_name,
	max_samples_per_dataset=max_samples,
	max_seq_len=32,
	per_device_batch_size=1,
	grad_accum=2,
	lr=1e-3,
	use_hypothesis=True,
	meta_interval=5,
	log_every=1,
	save_temp_every=100, # não salvar temp (modelo pequeno)
	output_dir="/home/z/my-project/BiGRU_T_version/model_final",
	temp_dir="/home/z/my-project/BiGRU_T_version/_temp_checkpoints",
	keep_temp=False,
	)

	trainer = BiGRU_T_Trainer(
	model=model,
	tokenizer=tokenizer,
	config=trainer_config,
	train_samples=train_samples,
	val_samples=val_samples,
	)

	result = trainer.train()

	print("\n=== Resultado Final ===")
	print(json.dumps(result, indent=2, default=str))

	# Apaga HF token
	if "HF_TOKEN" in os.environ:
	del os.environ["HF_TOKEN"]
	logger.info("HF_TOKEN apagado do ambiente")

	sys.exit(0 if not result.get("killed") else 1)


	if __name__ == "__main__":
	main()