|
| 1 | +#!/usr/bin/env python |
| 2 | +# coding: utf-8 |
| 3 | + |
| 4 | +# This script creates a super tiny model that is useful inside tests, when we just want to test that |
| 5 | +# the machinery works, without needing to the check the quality of the outcomes. |
| 6 | +# |
| 7 | +# This version creates a tiny vocab first, and then a tiny model - so the outcome is truly tiny - |
| 8 | +# all files ~60KB. As compared to taking a full-size model, reducing to the minimum its layers and |
| 9 | +# emb dimensions, but keeping the full vocab + merges files, leading to ~3MB in total for all files. |
| 10 | +# The latter is done by `fsmt-make-super-tiny-model.py`. |
| 11 | +# |
| 12 | +# It will be used then as "stas/tiny-wmt19-en-ru" |
| 13 | + |
| 14 | +from pathlib import Path |
| 15 | +import json |
| 16 | +import tempfile |
| 17 | + |
| 18 | +from transformers import FSMTTokenizer, FSMTConfig, FSMTForConditionalGeneration |
| 19 | +from transformers.tokenization_fsmt import VOCAB_FILES_NAMES |
| 20 | + |
| 21 | +mname_tiny = "tiny-wmt19-en-ru" |
| 22 | + |
| 23 | +# Build |
| 24 | + |
| 25 | +# borrowed from a test |
| 26 | +vocab = [ "l", "o", "w", "e", "r", "s", "t", "i", "d", "n", "w</w>", "r</w>", "t</w>", "lo", "low", "er</w>", "low</w>", "lowest</w>", "newer</w>", "wider</w>", "<unk>", ] |
| 27 | +vocab_tokens = dict(zip(vocab, range(len(vocab)))) |
| 28 | +merges = ["l o 123", "lo w 1456", "e r</w> 1789", ""] |
| 29 | + |
| 30 | +with tempfile.TemporaryDirectory() as tmpdirname: |
| 31 | + build_dir = Path(tmpdirname) |
| 32 | + src_vocab_file = build_dir / VOCAB_FILES_NAMES["src_vocab_file"] |
| 33 | + tgt_vocab_file = build_dir / VOCAB_FILES_NAMES["tgt_vocab_file"] |
| 34 | + merges_file = build_dir / VOCAB_FILES_NAMES["merges_file"] |
| 35 | + with open(src_vocab_file, "w") as fp: fp.write(json.dumps(vocab_tokens)) |
| 36 | + with open(tgt_vocab_file, "w") as fp: fp.write(json.dumps(vocab_tokens)) |
| 37 | + with open(merges_file, "w") as fp : fp.write("\n".join(merges)) |
| 38 | + |
| 39 | + tokenizer = FSMTTokenizer( |
| 40 | + langs=["en", "ru"], |
| 41 | + src_vocab_size = len(vocab), |
| 42 | + tgt_vocab_size = len(vocab), |
| 43 | + src_vocab_file=src_vocab_file, |
| 44 | + tgt_vocab_file=tgt_vocab_file, |
| 45 | + merges_file=merges_file, |
| 46 | + ) |
| 47 | + |
| 48 | +config = FSMTConfig( |
| 49 | + langs=['ru', 'en'], |
| 50 | + src_vocab_size=1000, tgt_vocab_size=1000, |
| 51 | + d_model=4, |
| 52 | + encoder_layers=1, decoder_layers=1, |
| 53 | + encoder_ffn_dim=4, decoder_ffn_dim=4, |
| 54 | + encoder_attention_heads=1, decoder_attention_heads=1, |
| 55 | +) |
| 56 | + |
| 57 | +tiny_model = FSMTForConditionalGeneration(config) |
| 58 | +print(f"num of params {tiny_model.num_parameters()}") |
| 59 | + |
| 60 | +# Test |
| 61 | +batch = tokenizer.prepare_seq2seq_batch(["Making tiny model"]) |
| 62 | +outputs = tiny_model(**batch, return_dict=True) |
| 63 | + |
| 64 | +print("test output:", len(outputs.logits[0])) |
| 65 | + |
| 66 | +# Save |
| 67 | +tiny_model.half() # makes it smaller |
| 68 | +tiny_model.save_pretrained(mname_tiny) |
| 69 | +tokenizer.save_pretrained(mname_tiny) |
| 70 | + |
| 71 | +print(f"Generated {mname_tiny}") |
| 72 | + |
| 73 | +# Upload |
| 74 | +# transformers-cli upload tiny-wmt19-en-ru |
0 commit comments