Spaces:
Sleeping
Sleeping
| #!/usr/bin/env python3 | |
| import os | |
| import torch | |
| from transformers import AutoModelForSeq2SeqLM, AutoModelForCausalLM | |
| from legacy.accel import SwitchingAccelerator | |
| from accelerate import Accelerator | |
| from data import MultilingualDatasetIterator | |
| from aux import log, CmdlineArgs | |
| from legacy.langconv import lang_set_maybe_smugri, is_dec_only_llm | |
| from modelops import mdl_param_count, to_cpl_spec, hf_tok | |
| from tokops import load_tokenizer | |
| def freeze_model(model): | |
| for n, p in model.named_parameters(): | |
| p.requires_grad = False | |
| def load_hf_mdl_and_tok(mdl_id, tok_id=None, verbose=False): | |
| if tok_id is None: | |
| tok_id = mdl_id | |
| tokenizer = load_tokenizer(tok_id) # AutoTokenizer.fromm_pretrained(tok_id, token=hf_tok) | |
| if is_dec_only_llm(tokenizer[0]): | |
| model = AutoModelForCausalLM.from_pretrained(mdl_id, token=hf_tok, torch_dtype=torch.bfloat16) | |
| else: | |
| model = AutoModelForSeq2SeqLM.from_pretrained(mdl_id, token=hf_tok, torch_dtype=torch.bfloat16) | |
| if verbose: | |
| mdl_size, _ = mdl_param_count(model) | |
| log(f"Loaded {mdl_id} with {mdl_size} params, voc size {model.config.vocab_size}") | |
| return model, tokenizer | |
| def _cmdline_args(): | |
| description = """Train or tune models""" | |
| pos_args = ["mdl_id", "save_location", "train_pretok_file", "langs"] | |
| pos_types = [str, str, str, lang_set_maybe_smugri] | |
| kw_args = { "anchor_mdl_id": None, "anchor_langs": None, "continue_training": False, | |
| "save_steps": 100000, "lr": 1.5e-5, "nr_snts_in_batch": 0, "nr_words_in_batch": 0, | |
| "log_steps": 100, "epochs": 4 } | |
| #post-process the arguments | |
| args = CmdlineArgs(description, pos_arg_list=pos_args, pos_arg_types=pos_types, kw_arg_dict=kw_args) | |
| if args.anchor_langs is not None: | |
| args.anchor_langs = lang_set_maybe_smugri(args.anchor_langs) | |
| if (args.nr_snts_in_batch > 0) == (args.nr_words_in_batch > 0): | |
| raise Exception(f"Specify the batch size either in words or in sentences.") | |
| # if the directory args.save_location already exists, raise an exception: | |
| if not args.continue_training and os.path.exists(args.save_location): | |
| raise Exception(f"Save location '{args.save_location}' already exists, don't want to overwrite.") | |
| return args | |
| def yes_i_called_this_function_do_main(): | |
| args = _cmdline_args() | |
| tmp_acc = Accelerator() | |
| log(f"Num proc: {tmp_acc.num_processes}, proc ID: {tmp_acc.process_index}") | |
| log("loading coupled model and tokenizer", accelerator=tmp_acc) | |
| main_model, main_tokenizer = load_hf_mdl_and_tok(args.mdl_id, verbose=True) | |
| coupling_specs = to_cpl_spec(args.langs, main_model, main_tokenizer[0], main_tokenizer[1], args.save_location) | |
| if args.anchor_mdl_id: | |
| log("loading anchor model and tokenizer", accelerator=tmp_acc) | |
| anchor_model, anchor_tokenizer = load_hf_mdl_and_tok(args.anchor_mdl_id, verbose=True) | |
| freeze_model(anchor_model) | |
| coupling_specs += to_cpl_spec(args.anchor_langs, anchor_model, anchor_tokenizer[0], anchor_tokenizer[1], args.anchor_mdl_id) | |
| train_set = MultilingualDatasetIterator(args.train_pretok_file) | |
| acc_trainer = SwitchingAccelerator(coupling_specs, train_set, args) | |
| upd_model, loss_list = acc_trainer.train() | |
| #save_all_models(args.save_location, upd_model, main_tokenizer, coupling_specs, loss_list, trainer=acc_trainer.accelerator) | |
| if __name__ == "__main__": | |
| #sys.argv = ". models/smol models/smol_next data/smugri4a-dev.json-tokcache/thiscache.json smugri log_steps=1 lr=1e-5".split() | |
| #sys.argv = ". models/llama3.2-1b models/llama-tuned data/smugri4a-dev.json-tokcache/llama.json smugri".split() | |
| yes_i_called_this_function_do_main() | |