hmnshudhmn24's picture
Upload 18 files
9a1ddeb verified
import yaml
from src.data.load_data import load_resume_dataset
from src.modeling.model import load_model
from src.modeling.tokenizer import load_tokenizer
from src.training.trainer import get_trainer
with open("src/config/model_config.yaml") as f:
model_cfg = yaml.safe_load(f)
with open("src/config/training_config.yaml") as f:
train_cfg = yaml.safe_load(f)
dataset = load_resume_dataset(
"data/processed/train.jsonl",
"data/processed/validation.jsonl",
"data/processed/test.jsonl"
)
tokenizer = load_tokenizer(model_cfg["model_name"])
model = load_model(
model_cfg["model_name"],
model_cfg["num_labels"],
model_cfg["id2label"],
model_cfg["label2id"]
)
trainer = get_trainer(model, tokenizer, dataset, train_cfg)
trainer.train()
trainer.save_model(train_cfg["output_dir"])