examples/finetune_vit5x_example.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8ffa2cd-7349-4320-b334-a25174b014f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -r ViT5\n",
    "!git clone https://github.com/vietAI/ViT5.git"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6568ab59-77ee-42e9-8142-5f4fdca89f33",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Install JAX for GPU\n",
    "!pip install jaxlib==0.4.2+cuda11.cudnn86 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\n",
    "## Install T5X and dependencies\n",
    "!cd ViT5 && python3 setup.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80277b11-d328-41b9-9416-47615cc8ecd0",
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_SIZE = \"base\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "522d162b-ce79-4f89-b585-e3ed94f58a64",
   "metadata": {},
   "outputs": [],
   "source": [
    "# download ViT5X base model\n",
    "!gsutil -m cp -r gs://vietai_public/viT5/ViT5_{MODEL_SIZE} ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af85ef06-c3c7-4824-a59c-90456ba28454",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "135d36e7-fbb5-42d7-9723-7027a6b2ad31",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "I0206 18:15:33.707216 140281490196288 finetune_t5x.py:623] Epoch 1000 of 1501\n",
      "I0206 18:15:33.707377 140281490196288 finetune_t5x.py:629] BEGIN Train loop.\n",
      "I0206 18:15:33.707416 140281490196288 finetune_t5x.py:634] Training for 1000 steps.\n",
      "I0206 18:15:33.708324 140223275026176 logging_writer.py:48] [1000000] collection=train timing/compilation_seconds=120.909\n",
      "I0206 18:15:33.710309 140223275026176 logging_writer.py:48] [1000000] collection=train timing/train_iter_warmup=6.19888e-06\n",
      "I0206 18:15:33.710495 140281490196288 trainer.py:500] Training: step 1000000\n",
      "I0206 18:15:43.957198 140281490196288 trainer.py:500] Training: step 1000025\n",
      "I0206 18:15:54.192222 140281490196288 trainer.py:500] Training: step 1000050\n",
      "I0206 18:16:04.432064 140281490196288 trainer.py:500] Training: step 1000075\n",
      "I0206 18:16:14.711035 140281490196288 trainer.py:500] Training: step 1000100\n",
      "I0206 18:16:24.777673 140281490196288 trainer.py:500] Training: step 1000124\n",
      "I0206 18:16:35.114796 140281490196288 trainer.py:500] Training: step 1000149\n",
      "I0206 18:16:45.492275 140281490196288 trainer.py:500] Training: step 1000174\n",
      "I0206 18:16:55.705092 140281490196288 trainer.py:500] Training: step 1000198\n",
      "I0206 18:17:05.707319 140281490196288 trainer.py:500] Training: step 1000222\n",
      "I0206 18:17:15.730436 140281490196288 trainer.py:500] Training: step 1000246\n",
      "I0206 18:17:25.765497 140281490196288 trainer.py:500] Training: step 1000270\n",
      "I0206 18:17:36.011804 140281490196288 trainer.py:500] Training: step 1000294\n",
      "I0206 18:17:46.052886 140281490196288 trainer.py:500] Training: step 1000318\n",
      "I0206 18:17:56.100627 140281490196288 trainer.py:500] Training: step 1000342\n",
      "I0206 18:18:06.150018 140281490196288 trainer.py:500] Training: step 1000366\n",
      "I0206 18:18:16.360258 140281490196288 trainer.py:500] Training: step 1000390\n",
      "I0206 18:18:26.421083 140281490196288 trainer.py:500] Training: step 1000414\n",
      "I0206 18:18:36.469063 140281490196288 trainer.py:500] Training: step 1000438\n",
      "I0206 18:18:46.521416 140281490196288 trainer.py:500] Training: step 1000462\n",
      "I0206 18:18:56.735538 140281490196288 trainer.py:500] Training: step 1000486\n",
      "I0206 18:19:06.789165 140281490196288 trainer.py:500] Training: step 1000510\n",
      "I0206 18:19:16.840288 140281490196288 trainer.py:500] Training: step 1000534\n",
      "I0206 18:19:26.896516 140281490196288 trainer.py:500] Training: step 1000558\n",
      "I0206 18:19:37.104189 140281490196288 trainer.py:500] Training: step 1000582\n",
      "I0206 18:19:47.156475 140281490196288 trainer.py:500] Training: step 1000606\n",
      "I0206 18:19:57.209078 140281490196288 trainer.py:500] Training: step 1000630\n",
      "I0206 18:20:07.263255 140281490196288 trainer.py:500] Training: step 1000654\n",
      "I0206 18:20:17.506186 140281490196288 trainer.py:500] Training: step 1000678\n",
      "I0206 18:20:27.556405 140281490196288 trainer.py:500] Training: step 1000702\n",
      "I0206 18:20:37.612117 140281490196288 trainer.py:500] Training: step 1000726\n",
      "I0206 18:20:47.666839 140281490196288 trainer.py:500] Training: step 1000750\n",
      "I0206 18:20:57.923375 140281490196288 trainer.py:500] Training: step 1000774\n",
      "I0206 18:21:07.978493 140281490196288 trainer.py:500] Training: step 1000798\n",
      "I0206 18:21:18.031299 140281490196288 trainer.py:500] Training: step 1000822\n",
      "I0206 18:21:28.085952 140281490196288 trainer.py:500] Training: step 1000846\n",
      "I0206 18:21:38.304044 140281490196288 trainer.py:500] Training: step 1000870\n",
      "I0206 18:21:48.361248 140281490196288 trainer.py:500] Training: step 1000894\n",
      "I0206 18:21:58.413504 140281490196288 trainer.py:500] Training: step 1000918\n",
      "I0206 18:22:08.620894 140281490196288 trainer.py:500] Training: step 1000942\n",
      "I0206 18:22:18.673136 140281490196288 trainer.py:500] Training: step 1000966\n",
      "I0206 18:22:28.723974 140281490196288 trainer.py:500] Training: step 1000990\n",
      "I0206 18:22:33.135376 140281490196288 finetune_t5x.py:659] END Train loop.\n",
      "I0206 18:22:33.135539 140281490196288 finetune_t5x.py:498] Running inference evaluation.\n",
      "I0206 18:22:33.145542 140281490196288 evaluation.py:628] Evaluating FAQ_summarization\n",
      "I0206 18:22:33.155588 140223275026176 logging_writer.py:48] [1001000] collection=train accuracy=0.816739, cross_ent_loss=1157.55, cross_ent_loss_per_all_target_tokens=0.282606, effective_batch_size/decoder=115.484, effective_batch_size/encoder=115.484, learning_rate=0.000999991, learning_rate/current=0.0010000000474974513, loss=1160.06, loss_per_all_target_tokens=0.283217, loss_per_nonpadding_target_token=0.762185, non_padding_fraction/decoder=0.371585, non_padding_fraction/encoder=0.773758, non_padding_fraction/loss_weights=0.371585, non_padding_fraction/overall=0.693323, timing/seconds=419.403, timing/seqs=64000, timing/seqs_per_second=152.598, timing/seqs_per_second_per_core=152.598, timing/steps_per_second=2.38434, timing/target_tokens_per_second=9766.27, timing/target_tokens_per_second_per_core=9766.27, timing/uptime=544.9464721679688, z_loss=2.50461, z_loss_per_all_target_tokens=0.000611476\n",
      "I0206 18:22:33.343189 140281490196288 utils.py:1329] length of dataset = 1330\n",
      "I0206 18:22:33.343356 140281490196288 utils.py:1340] Padding infer dataset with 14 examples for even per-replica shards.\n",
      "I0206 18:22:33.461422 140281490196288 utils.py:1355] The infer dataset is sharded into 1 shards with per-shard batch size of 64\n",
      "I0206 18:22:52.953899 140281490196288 utils.py:1383] Inference of batch [ 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23\n",
      " 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47\n",
      " 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63] done.\n",
      "I0206 18:22:56.856373 140281490196288 utils.py:1383] Inference of batch [ 64  65  66  67  68  69  70  71  72  73  74  75  76  77  78  79  80  81\n",
      "  82  83  84  85  86  87  88  89  90  91  92  93  94  95  96  97  98  99\n",
      " 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117\n",
      " 118 119 120 121 122 123 124 125 126 127] done.\n",
      "I0206 18:23:00.169296 140281490196288 utils.py:1383] Inference of batch [128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145\n",
      " 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163\n",
      " 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181\n",
      " 182 183 184 185 186 187 188 189 190 191] done.\n",
      "I0206 18:23:03.880506 140281490196288 utils.py:1383] Inference of batch [192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209\n",
      " 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227\n",
      " 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245\n",
      " 246 247 248 249 250 251 252 253 254 255] done.\n",
      "I0206 18:23:07.204680 140281490196288 utils.py:1383] Inference of batch [256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273\n",
      " 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291\n",
      " 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309\n",
      " 310 311 312 313 314 315 316 317 318 319] done.\n",
      "I0206 18:23:10.920166 140281490196288 utils.py:1383] Inference of batch [320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337\n",
      " 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355\n",
      " 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373\n",
      " 374 375 376 377 378 379 380 381 382 383] done.\n",
      "I0206 18:23:14.245954 140281490196288 utils.py:1383] Inference of batch [384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401\n",
      " 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419\n",
      " 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437\n",
      " 438 439 440 441 442 443 444 445 446 447] done.\n",
      "I0206 18:23:17.967159 140281490196288 utils.py:1383] Inference of batch [448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465\n",
      " 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483\n",
      " 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501\n",
      " 502 503 504 505 506 507 508 509 510 511] done.\n",
      "I0206 18:23:21.197070 140281490196288 utils.py:1383] Inference of batch [512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529\n",
      " 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547\n",
      " 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565\n",
      " 566 567 568 569 570 571 572 573 574 575] done.\n",
      "I0206 18:23:24.231013 140281490196288 utils.py:1383] Inference of batch [576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593\n",
      " 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611\n",
      " 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629\n",
      " 630 631 632 633 634 635 636 637 638 639] done.\n",
      "I0206 18:23:28.158111 140281490196288 utils.py:1383] Inference of batch [640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657\n",
      " 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675\n",
      " 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693\n",
      " 694 695 696 697 698 699 700 701 702 703] done.\n",
      "I0206 18:23:31.692280 140281490196288 utils.py:1383] Inference of batch [704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721\n",
      " 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739\n",
      " 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757\n",
      " 758 759 760 761 762 763 764 765 766 767] done.\n",
      "I0206 18:23:35.621019 140281490196288 utils.py:1383] Inference of batch [768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785\n",
      " 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803\n",
      " 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821\n",
      " 822 823 824 825 826 827 828 829 830 831] done.\n",
      "I0206 18:23:39.251537 140281490196288 utils.py:1383] Inference of batch [832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849\n",
      " 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867\n",
      " 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885\n",
      " 886 887 888 889 890 891 892 893 894 895] done.\n",
      "I0206 18:23:43.081492 140281490196288 utils.py:1383] Inference of batch [896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913\n",
      " 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931\n",
      " 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949\n",
      " 950 951 952 953 954 955 956 957 958 959] done.\n",
      "I0206 18:23:46.716987 140281490196288 utils.py:1383] Inference of batch [ 960  961  962  963  964  965  966  967  968  969  970  971  972  973\n",
      "  974  975  976  977  978  979  980  981  982  983  984  985  986  987\n",
      "  988  989  990  991  992  993  994  995  996  997  998  999 1000 1001\n",
      " 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015\n",
      " 1016 1017 1018 1019 1020 1021 1022 1023] done.\n",
      "I0206 18:23:50.058820 140281490196288 utils.py:1383] Inference of batch [1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037\n",
      " 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051\n",
      " 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065\n",
      " 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079\n",
      " 1080 1081 1082 1083 1084 1085 1086 1087] done.\n",
      "I0206 18:23:53.501441 140281490196288 utils.py:1383] Inference of batch [1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101\n",
      " 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115\n",
      " 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126 1127 1128 1129\n",
      " 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143\n",
      " 1144 1145 1146 1147 1148 1149 1150 1151] done.\n",
      "I0206 18:23:57.442556 140281490196288 utils.py:1383] Inference of batch [1152 1153 1154 1155 1156 1157 1158 1159 1160 1161 1162 1163 1164 1165\n",
      " 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179\n",
      " 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193\n",
      " 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207\n",
      " 1208 1209 1210 1211 1212 1213 1214 1215] done.\n",
      "I0206 18:24:00.584612 140281490196288 utils.py:1383] Inference of batch [1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229\n",
      " 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243\n",
      " 1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257\n",
      " 1258 1259 1260 1261 1262 1263 1264 1265 1266 1267 1268 1269 1270 1271\n",
      " 1272 1273 1274 1275 1276 1277 1278 1279] done.\n",
      "I0206 18:24:03.521842 140281490196288 utils.py:1383] Inference of batch [1280 1281 1282 1283 1284 1285 1286 1287 1288 1289 1290 1291 1292 1293\n",
      " 1294 1295 1296 1297 1298 1299 1300 1301 1302 1303 1304 1305 1306 1307\n",
      " 1308 1309 1310 1311 1312 1313 1314 1315 1316 1317 1318 1319 1320 1321\n",
      " 1322 1323 1324 1325 1326 1327 1328 1329   -1   -1   -1   -1   -1   -1\n",
      "   -1   -1   -1   -1   -1   -1   -1   -1] done.\n",
      "I0206 18:24:03.526962 140281490196288 utils.py:1408] Inference of all batches done.\n",
      "I0206 18:24:03.539730 140223987750656 evaluation.py:698] Computing metrics for FAQ_summarization\n",
      "I0206 18:24:03.539989 140281490196288 finetune_t5x.py:623] Epoch 1001 of 1501\n",
      "I0206 18:24:03.540841 140223275026176 logging_writer.py:48] [1001000] collection=train timing/evaluate_seconds=90.4043\n",
      "I0206 18:24:03.541774 140281490196288 finetune_t5x.py:629] BEGIN Train loop.\n",
      "I0206 18:24:03.542388 140281490196288 finetune_t5x.py:634] Training for 1000 steps.\n",
      "I0206 18:24:03.542513 140281490196288 trainer.py:500] Training: step 1001000\n",
      "Downloading builder script: 5.60kB [00:00, 5.01MB/s]                            \n",
      "I0206 18:24:04.491727 140223987750656 rouge_scorer.py:83] Using default tokenizer.\n",
      "I0206 18:24:05.140494 140223987750656 loggers.py:105] FAQ_summarization/rouge1 at step 1001000: 0.663\n",
      "I0206 18:24:05.140709 140223987750656 loggers.py:105] FAQ_summarization/rouge2 at step 1001000: 0.516\n",
      "I0206 18:24:05.140753 140223987750656 loggers.py:105] FAQ_summarization/rougeL at step 1001000: 0.611\n",
      "I0206 18:24:05.140787 140223987750656 loggers.py:105] FAQ_summarization/rougeLsum at step 1001000: 0.612\n",
      "I0206 18:24:05.167299 140223987750656 loggers.py:426] Appending metrics to out/FAQ_summarization/vit5_base/inference_eval/FAQ_summarization-metrics.jsonl\n",
      "I0206 18:24:05.168845 140223987750656 loggers.py:458] Writing inferences to out/FAQ_summarization/vit5_base/inference_eval/FAQ_summarization-1001000.jsonl\n",
      "I0206 18:24:05.513097 140223987750656 loggers.py:500] Writing completed in 0.344246 seconds (5.809804 examples/sec).\n",
      "I0206 18:24:05.513332 140223987750656 evaluation.py:655] Time computing metrics: 1.973633 secs.\n",
      "I0206 18:24:13.883319 140281490196288 trainer.py:500] Training: step 1001025\n",
      "I0206 18:24:23.917412 140281490196288 trainer.py:500] Training: step 1001049\n",
      "I0206 18:24:34.202437 140281490196288 trainer.py:500] Training: step 1001073\n",
      "I0206 18:24:44.233891 140281490196288 trainer.py:500] Training: step 1001097\n",
      "I0206 18:24:54.270677 140281490196288 trainer.py:500] Training: step 1001121\n",
      "I0206 18:25:04.308255 140281490196288 trainer.py:500] Training: step 1001145\n",
      "I0206 18:25:14.345350 140281490196288 trainer.py:500] Training: step 1001169\n",
      "I0206 18:25:24.579156 140281490196288 trainer.py:500] Training: step 1001193\n",
      "I0206 18:25:34.615727 140281490196288 trainer.py:500] Training: step 1001217\n",
      "I0206 18:25:44.649321 140281490196288 trainer.py:500] Training: step 1001241\n",
      "I0206 18:25:54.689231 140281490196288 trainer.py:500] Training: step 1001265\n",
      "I0206 18:26:04.902699 140281490196288 trainer.py:500] Training: step 1001289\n",
      "I0206 18:26:14.932554 140281490196288 trainer.py:500] Training: step 1001313\n",
      "I0206 18:26:24.964614 140281490196288 trainer.py:500] Training: step 1001337\n",
      "I0206 18:26:35.002245 140281490196288 trainer.py:500] Training: step 1001361\n",
      "I0206 18:26:45.234210 140281490196288 trainer.py:500] Training: step 1001385\n",
      "I0206 18:26:55.267924 140281490196288 trainer.py:500] Training: step 1001409\n",
      "I0206 18:27:05.299423 140281490196288 trainer.py:500] Training: step 1001433\n",
      "I0206 18:27:15.333166 140281490196288 trainer.py:500] Training: step 1001457\n",
      "I0206 18:27:25.584102 140281490196288 trainer.py:500] Training: step 1001481\n",
      "I0206 18:27:35.612276 140281490196288 trainer.py:500] Training: step 1001505\n",
      "I0206 18:27:45.644881 140281490196288 trainer.py:500] Training: step 1001529\n",
      "I0206 18:27:55.670200 140281490196288 trainer.py:500] Training: step 1001553\n",
      "I0206 18:28:05.940177 140281490196288 trainer.py:500] Training: step 1001577\n",
      "I0206 18:28:15.966579 140281490196288 trainer.py:500] Training: step 1001601\n",
      "I0206 18:28:25.997464 140281490196288 trainer.py:500] Training: step 1001625\n",
      "I0206 18:28:36.023604 140281490196288 trainer.py:500] Training: step 1001649\n",
      "I0206 18:28:46.287488 140281490196288 trainer.py:500] Training: step 1001673\n",
      "I0206 18:28:56.310695 140281490196288 trainer.py:500] Training: step 1001697\n",
      "I0206 18:29:06.336511 140281490196288 trainer.py:500] Training: step 1001721\n",
      "I0206 18:29:16.362411 140281490196288 trainer.py:500] Training: step 1001745\n",
      "I0206 18:29:26.590661 140281490196288 trainer.py:500] Training: step 1001769\n",
      "I0206 18:29:36.616428 140281490196288 trainer.py:500] Training: step 1001793\n",
      "I0206 18:29:46.642745 140281490196288 trainer.py:500] Training: step 1001817\n",
      "I0206 18:29:56.668151 140281490196288 trainer.py:500] Training: step 1001841\n",
      "I0206 18:30:06.694212 140281490196288 trainer.py:500] Training: step 1001865\n",
      "I0206 18:30:16.981838 140281490196288 trainer.py:500] Training: step 1001889\n",
      "I0206 18:30:27.007698 140281490196288 trainer.py:500] Training: step 1001913\n",
      "I0206 18:30:37.030871 140281490196288 trainer.py:500] Training: step 1001937\n",
      "I0206 18:30:47.055918 140281490196288 trainer.py:500] Training: step 1001961\n",
      "I0206 18:30:57.322506 140281490196288 trainer.py:500] Training: step 1001985\n",
      "I0206 18:31:03.592247 140281490196288 finetune_t5x.py:659] END Train loop.\n",
      "I0206 18:31:03.592403 140281490196288 finetune_t5x.py:498] Running inference evaluation.\n",
      "I0206 18:31:03.716535 140281490196288 evaluation.py:628] Evaluating FAQ_summarization\n",
      "I0206 18:31:03.722774 140281490196288 utils.py:1329] length of dataset = 1330\n",
      "I0206 18:31:03.722888 140281490196288 utils.py:1340] Padding infer dataset with 14 examples for even per-replica shards.\n",
      "I0206 18:31:03.735070 140223275026176 logging_writer.py:48] [1002000] collection=train accuracy=0.926701, cross_ent_loss=366.069, cross_ent_loss_per_all_target_tokens=0.0893724, effective_batch_size/decoder=115.449, effective_batch_size/encoder=115.449, learning_rate=0.000999991, learning_rate/current=0.0010000000474974513, loss=369.87, loss_per_all_target_tokens=0.0903003, loss_per_nonpadding_target_token=0.243071, non_padding_fraction/decoder=0.371497, non_padding_fraction/encoder=0.773533, non_padding_fraction/loss_weights=0.371497, non_padding_fraction/overall=0.693126, timing/seconds=420.158, timing/seqs=64000, timing/seqs_per_second=152.324, timing/seqs_per_second_per_core=152.324, timing/steps_per_second=2.38006, timing/target_tokens_per_second=9748.71, timing/target_tokens_per_second_per_core=9748.71, timing/uptime=1055.6163330078125, z_loss=3.80059, z_loss_per_all_target_tokens=0.00092788\n",
      "I0206 18:31:03.742032 140281490196288 utils.py:1355] The infer dataset is sharded into 1 shards with per-shard batch size of 64\n",
      "I0206 18:31:07.111604 140281490196288 utils.py:1383] Inference of batch [ 0  1  2  3  4  5  6  7  8  9 10 11 12 13 14 15 16 17 18 19 20 21 22 23\n",
      " 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47\n",
      " 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63] done.\n",
      "I0206 18:31:11.047202 140281490196288 utils.py:1383] Inference of batch [ 64  65  66  67  68  69  70  71  72  73  74  75  76  77  78  79  80  81\n",
      "  82  83  84  85  86  87  88  89  90  91  92  93  94  95  96  97  98  99\n",
      " 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117\n",
      " 118 119 120 121 122 123 124 125 126 127] done.\n",
      "I0206 18:31:14.387603 140281490196288 utils.py:1383] Inference of batch [128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145\n",
      " 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163\n",
      " 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181\n",
      " 182 183 184 185 186 187 188 189 190 191] done.\n",
      "I0206 18:31:18.123245 140281490196288 utils.py:1383] Inference of batch [192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209\n",
      " 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227\n",
      " 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245\n",
      " 246 247 248 249 250 251 252 253 254 255] done.\n",
      "I0206 18:31:20.863952 140281490196288 utils.py:1383] Inference of batch [256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273\n",
      " 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291\n",
      " 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 309\n",
      " 310 311 312 313 314 315 316 317 318 319] done.\n",
      "I0206 18:31:24.399050 140281490196288 utils.py:1383] Inference of batch [320 321 322 323 324 325 326 327 328 329 330 331 332 333 334 335 336 337\n",
      " 338 339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355\n",
      " 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371 372 373\n",
      " 374 375 376 377 378 379 380 381 382 383] done.\n",
      "I0206 18:31:27.540469 140281490196288 utils.py:1383] Inference of batch [384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401\n",
      " 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419\n",
      " 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437\n",
      " 438 439 440 441 442 443 444 445 446 447] done.\n",
      "I0206 18:31:30.480320 140281490196288 utils.py:1383] Inference of batch [448 449 450 451 452 453 454 455 456 457 458 459 460 461 462 463 464 465\n",
      " 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483\n",
      " 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501\n",
      " 502 503 504 505 506 507 508 509 510 511] done.\n",
      "I0206 18:31:34.018755 140281490196288 utils.py:1383] Inference of batch [512 513 514 515 516 517 518 519 520 521 522 523 524 525 526 527 528 529\n",
      " 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547\n",
      " 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565\n",
      " 566 567 568 569 570 571 572 573 574 575] done.\n",
      "I0206 18:31:37.367018 140281490196288 utils.py:1383] Inference of batch [576 577 578 579 580 581 582 583 584 585 586 587 588 589 590 591 592 593\n",
      " 594 595 596 597 598 599 600 601 602 603 604 605 606 607 608 609 610 611\n",
      " 612 613 614 615 616 617 618 619 620 621 622 623 624 625 626 627 628 629\n",
      " 630 631 632 633 634 635 636 637 638 639] done.\n",
      "I0206 18:31:40.911081 140281490196288 utils.py:1383] Inference of batch [640 641 642 643 644 645 646 647 648 649 650 651 652 653 654 655 656 657\n",
      " 658 659 660 661 662 663 664 665 666 667 668 669 670 671 672 673 674 675\n",
      " 676 677 678 679 680 681 682 683 684 685 686 687 688 689 690 691 692 693\n",
      " 694 695 696 697 698 699 700 701 702 703] done.\n",
      "I0206 18:31:44.352218 140281490196288 utils.py:1383] Inference of batch [704 705 706 707 708 709 710 711 712 713 714 715 716 717 718 719 720 721\n",
      " 722 723 724 725 726 727 728 729 730 731 732 733 734 735 736 737 738 739\n",
      " 740 741 742 743 744 745 746 747 748 749 750 751 752 753 754 755 756 757\n",
      " 758 759 760 761 762 763 764 765 766 767] done.\n",
      "I0206 18:31:47.796376 140281490196288 utils.py:1383] Inference of batch [768 769 770 771 772 773 774 775 776 777 778 779 780 781 782 783 784 785\n",
      " 786 787 788 789 790 791 792 793 794 795 796 797 798 799 800 801 802 803\n",
      " 804 805 806 807 808 809 810 811 812 813 814 815 816 817 818 819 820 821\n",
      " 822 823 824 825 826 827 828 829 830 831] done.\n",
      "I0206 18:31:51.435350 140281490196288 utils.py:1383] Inference of batch [832 833 834 835 836 837 838 839 840 841 842 843 844 845 846 847 848 849\n",
      " 850 851 852 853 854 855 856 857 858 859 860 861 862 863 864 865 866 867\n",
      " 868 869 870 871 872 873 874 875 876 877 878 879 880 881 882 883 884 885\n",
      " 886 887 888 889 890 891 892 893 894 895] done.\n",
      "I0206 18:31:55.172258 140281490196288 utils.py:1383] Inference of batch [896 897 898 899 900 901 902 903 904 905 906 907 908 909 910 911 912 913\n",
      " 914 915 916 917 918 919 920 921 922 923 924 925 926 927 928 929 930 931\n",
      " 932 933 934 935 936 937 938 939 940 941 942 943 944 945 946 947 948 949\n",
      " 950 951 952 953 954 955 956 957 958 959] done.\n",
      "I0206 18:31:58.510827 140281490196288 utils.py:1383] Inference of batch [ 960  961  962  963  964  965  966  967  968  969  970  971  972  973\n",
      "  974  975  976  977  978  979  980  981  982  983  984  985  986  987\n",
      "  988  989  990  991  992  993  994  995  996  997  998  999 1000 1001\n",
      " 1002 1003 1004 1005 1006 1007 1008 1009 1010 1011 1012 1013 1014 1015\n",
      " 1016 1017 1018 1019 1020 1021 1022 1023] done.\n",
      "I0206 18:32:02.259103 140281490196288 utils.py:1383] Inference of batch [1024 1025 1026 1027 1028 1029 1030 1031 1032 1033 1034 1035 1036 1037\n",
      " 1038 1039 1040 1041 1042 1043 1044 1045 1046 1047 1048 1049 1050 1051\n",
      " 1052 1053 1054 1055 1056 1057 1058 1059 1060 1061 1062 1063 1064 1065\n",
      " 1066 1067 1068 1069 1070 1071 1072 1073 1074 1075 1076 1077 1078 1079\n",
      " 1080 1081 1082 1083 1084 1085 1086 1087] done.\n",
      "I0206 18:32:05.407696 140281490196288 utils.py:1383] Inference of batch [1088 1089 1090 1091 1092 1093 1094 1095 1096 1097 1098 1099 1100 1101\n",
      " 1102 1103 1104 1105 1106 1107 1108 1109 1110 1111 1112 1113 1114 1115\n",
      " 1116 1117 1118 1119 1120 1121 1122 1123 1124 1125 1126 1127 1128 1129\n",
      " 1130 1131 1132 1133 1134 1135 1136 1137 1138 1139 1140 1141 1142 1143\n",
      " 1144 1145 1146 1147 1148 1149 1150 1151] done.\n",
      "I0206 18:32:08.355411 140281490196288 utils.py:1383] Inference of batch [1152 1153 1154 1155 1156 1157 1158 1159 1160 1161 1162 1163 1164 1165\n",
      " 1166 1167 1168 1169 1170 1171 1172 1173 1174 1175 1176 1177 1178 1179\n",
      " 1180 1181 1182 1183 1184 1185 1186 1187 1188 1189 1190 1191 1192 1193\n",
      " 1194 1195 1196 1197 1198 1199 1200 1201 1202 1203 1204 1205 1206 1207\n",
      " 1208 1209 1210 1211 1212 1213 1214 1215] done.\n",
      "I0206 18:32:11.501761 140281490196288 utils.py:1383] Inference of batch [1216 1217 1218 1219 1220 1221 1222 1223 1224 1225 1226 1227 1228 1229\n",
      " 1230 1231 1232 1233 1234 1235 1236 1237 1238 1239 1240 1241 1242 1243\n",
      " 1244 1245 1246 1247 1248 1249 1250 1251 1252 1253 1254 1255 1256 1257\n",
      " 1258 1259 1260 1261 1262 1263 1264 1265 1266 1267 1268 1269 1270 1271\n",
      " 1272 1273 1274 1275 1276 1277 1278 1279] done.\n",
      "I0206 18:32:14.441715 140281490196288 utils.py:1383] Inference of batch [1280 1281 1282 1283 1284 1285 1286 1287 1288 1289 1290 1291 1292 1293\n",
      " 1294 1295 1296 1297 1298 1299 1300 1301 1302 1303 1304 1305 1306 1307\n",
      " 1308 1309 1310 1311 1312 1313 1314 1315 1316 1317 1318 1319 1320 1321\n",
      " 1322 1323 1324 1325 1326 1327 1328 1329   -1   -1   -1   -1   -1   -1\n",
      "   -1   -1   -1   -1   -1   -1   -1   -1] done.\n",
      "I0206 18:32:25.882165 140281490196288 utils.py:1408] Inference of all batches done.\n",
      "I0206 18:32:25.897804 140281490196288 evaluation.py:647] Time waiting for previous metrics run: 0.000021 secs.\n",
      "I0206 18:32:25.898162 140281490196288 finetune_t5x.py:623] Epoch 1002 of 1501\n",
      "I0206 18:32:25.898257 140281490196288 finetune_t5x.py:629] BEGIN Train loop.\n",
      "I0206 18:32:25.898301 140281490196288 finetune_t5x.py:634] Training for 1000 steps.\n",
      "I0206 18:32:25.898391 140281490196288 trainer.py:500] Training: step 1002000\n",
      "I0206 18:32:25.898758 140223275026176 logging_writer.py:48] [1002000] collection=train timing/evaluate_seconds=82.3056\n",
      "I0206 18:32:25.899148 140223987750656 evaluation.py:698] Computing metrics for FAQ_summarization\n",
      "I0206 18:32:26.388127 140223987750656 rouge_scorer.py:83] Using default tokenizer.\n",
      "I0206 18:32:27.027732 140223987750656 loggers.py:105] FAQ_summarization/rouge1 at step 1002000: 0.654\n",
      "I0206 18:32:27.027891 140223987750656 loggers.py:105] FAQ_summarization/rouge2 at step 1002000: 0.507\n",
      "I0206 18:32:27.027930 140223987750656 loggers.py:105] FAQ_summarization/rougeL at step 1002000: 0.601\n",
      "I0206 18:32:27.027965 140223987750656 loggers.py:105] FAQ_summarization/rougeLsum at step 1002000: 0.601\n",
      "I0206 18:32:27.030178 140223987750656 loggers.py:426] Appending metrics to out/FAQ_summarization/vit5_base/inference_eval/FAQ_summarization-metrics.jsonl\n",
      "I0206 18:32:27.066222 140223987750656 loggers.py:458] Writing inferences to out/FAQ_summarization/vit5_base/inference_eval/FAQ_summarization-1002000.jsonl\n",
      "I0206 18:32:27.417734 140223987750656 loggers.py:500] Writing completed in 0.351532 seconds (5.689387 examples/sec).\n",
      "I0206 18:32:27.417887 140223987750656 evaluation.py:655] Time computing metrics: 1.518780 secs.\n",
      "I0206 18:32:36.199214 140281490196288 trainer.py:500] Training: step 1002025\n"
     ]
    }
   ],
   "source": [
    "############################### Generation Task ###############################\n",
    "######################### A Health-domain Summarization Dataset released by ######################\n",
    "######################### ViHealthBERT (https://aclanthology.org/2022.lrec-1.35/) #################\n",
    "\n",
    "!mkdir FAQ_summarization\n",
    "!wget -O FAQ_summarization/dev.tsv https://raw.githubusercontent.com/vietai/ViPubmed/main/data/FAQ_summarization/dev.tsv\n",
    "!wget -O FAQ_summarization/test.tsv https://raw.githubusercontent.com/vietai/ViPubmed/main/data/FAQ_summarization/test.tsv\n",
    "!wget -O FAQ_summarization/train.tsv https://raw.githubusercontent.com/vietai/ViPubmed/main/data/FAQ_summarization/train.tsv\n",
    "\n",
    "################################################################# \n",
    "\n",
    "task = 'FAQ_summarization'\n",
    "train_file = f'{task}/train.tsv'\n",
    "test_file = f'{task}/test.tsv'\n",
    "dev_file = f'{task}/dev.tsv'\n",
    "\n",
    "model_dir = f'out/{task}/vit5_base'\n",
    "pretrained_path=f'ViT5_{MODEL_SIZE}/checkpoint_1000000'\n",
    "\n",
    "gin_file = f'ViT5/configs/runs/{MODEL_SIZE}_finetune.gin'\n",
    "\n",
    "metric = 'rouge'\n",
    "\n",
    "# Train settings\n",
    "batch_size = 64\n",
    "features_length = {\"inputs\": 256, \"targets\": 64}\n",
    "train_steps = 1000 + 1500000 # 1000 finetune steps + 1.5M pretraining step\n",
    "save_period = 1000\n",
    "eval_period = 1000\n",
    "\n",
    "!python3 'ViT5/src/finetune_t5x.py' \\\n",
    "  --gin_file=\"{gin_file}\" \\\n",
    "  --gin.MODEL_DIR=\"'{model_dir}'\" \\\n",
    "  --gin.INITIAL_CHECKPOINT_PATH=\"'{pretrained_path}'\" \\\n",
    "  --gin.TRAIN_STEPS='{train_steps}' \\\n",
    "  --gin.SAVE_PERIOD='{save_period}'\\\n",
    "  --gin.EVAL_PERIOD='{eval_period}'\\\n",
    "  --gin.MIXTURE_OR_TASK_NAME=\"'{task}'\" \\\n",
    "  --gin.TASK_FEATURE_LENGTHS=\"{features_length}\" \\\n",
    "  --gin.BATCH_SIZE='{batch_size}' \\\n",
    "  --task=\"{task}\" \\\n",
    "  --metric=\"{metric}\" \\\n",
    "  --train_file=\"{train_file}\" \\\n",
    "  --predict_file=\"{test_file}\" # or {dev_file}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f783b95-3880-4860-ac20-eba207d3a369",
   "metadata": {},
   "outputs": [],
   "source": [
    "############################### Classification Task ##########################\n",
    "############################### ViMedNLI #####################################\n",
    "############## https://arxiv.org/abs/2210.05610 ##############################\n",
    "\n",
    "!mkdir vi_mednli\n",
    "!wget -O vi_mednli/dev.tsv https://raw.githubusercontent.com/vietai/ViPubmed/main/data/vi_mednli/dev_vi_refined.tsv \n",
    "!wget -O vi_mednli/test.tsv https://raw.githubusercontent.com/vietai/ViPubmed/main/data/vi_mednli/test_vi_refined.tsv \n",
    "!wget -O vi_mednli/train.tsv https://raw.githubusercontent.com/vietai/ViPubmed/main/data/vi_mednli/train_vi_refined.tsv \n",
    "\n",
    "############################### vimednli ###############################\n",
    "\n",
    "MODEL_SIZE = \"base\"\n",
    "task = 'vi_mednli'\n",
    "train_file = f'{task}/train.tsv'\n",
    "test_file = f'{task}/test.tsv'\n",
    "dev_file = f'{task}/dev.tsv'\n",
    "\n",
    "model_dir = f'out/{task}/vit5_base'\n",
    "pretrained_path=f'ViT5_{MODEL_SIZE}/checkpoint_1000000'\n",
    "\n",
    "gin_file = f'ViT5/configs/runs/{MODEL_SIZE}_finetune.gin'\n",
    "\n",
    "metric = 'accuracy'\n",
    "\n",
    "# Train settings\n",
    "batch_size = 64\n",
    "features_length = {\"inputs\": 128, \"targets\": 6}\n",
    "train_steps = 4800 + 1000000 # 1000 finetune steps + 1.5M pretraining step\n",
    "save_period = 4800\n",
    "eval_period = 4800\n",
    "learning_rate = 0.0005\n",
    "\n",
    "!python3 'ViT5/src/finetune_t5x.py' \\\n",
    "  --gin_file=\"{gin_file}\" \\\n",
    "  --gin.MODEL_DIR=\"'{model_dir}'\" \\\n",
    "  --gin.INITIAL_CHECKPOINT_PATH=\"'{pretrained_path}'\" \\\n",
    "  --gin.TRAIN_STEPS='{train_steps}' \\\n",
    "  --gin.SAVE_PERIOD='{save_period}'\\\n",
    "  --gin.EVAL_PERIOD='{eval_period}'\\\n",
    "  --gin.MIXTURE_OR_TASK_NAME=\"'{task}'\" \\\n",
    "  --gin.LEARNING_RATE='{learning_rate}' \\\n",
    "  --gin.TASK_FEATURE_LENGTHS=\"{features_length}\" \\\n",
    "  --gin.BATCH_SIZE='{batch_size}' \\\n",
    "  --task=\"{task}\" \\\n",
    "  --metric=\"{metric}\" \\\n",
    "  --train_file=\"{train_file}\" \\\n",
    "  --predict_file=\"{test_file}\" # or {dev_file}\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3aac4fd3-05c3-4269-b314-4632d1949140",
   "metadata": {},
   "outputs": [],
   "source": [
    "# download vietnews or wikilingua dataset\n",
    "!gsutil cp -r gs://vietai_public/viT5/data/wikilingua .\n",
    "# !gsutil cp -r gs://vietai_public/viT5/data/vietnews ."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6b80d14-21d3-4d61-a1fa-b636cd59e027",
   "metadata": {},
   "outputs": [],
   "source": [
    "############################### Generation Task ##########################\n",
    "############################### wikilingua ###############################\n",
    "MODEL_SIZE = \"base\"\n",
    "task = 'wikilingua'\n",
    "train_file = f'{task}/train.tsv'\n",
    "test_file = f'{task}/test.tsv'\n",
    "dev_file = f'{task}/dev.tsv'\n",
    "\n",
    "model_dir = f'out/{task}/vit5_{MODEL_SIZE}'\n",
    "pretrained_path=f'ViT5_{MODEL_SIZE}/checkpoint_1000000'\n",
    "\n",
    "gin_file = f'ViT5/configs/runs/{MODEL_SIZE}_finetune.gin'\n",
    "metric = 'rouge'\n",
    "\n",
    "# Train settings\n",
    "batch_size = 16\n",
    "features_length = {\"inputs\": 1024, \"targets\": 256}\n",
    "train_steps = 10000 + 1000000 # 1000 finetune steps + 1.5M pretraining step\n",
    "save_period = 10000\n",
    "eval_period = 10000\n",
    "\n",
    "!python3 'ViT5/src/finetune_t5x.py' \\\n",
    "  --gin_file=\"{gin_file}\" \\\n",
    "  --gin.MODEL_DIR=\"'{model_dir}'\" \\\n",
    "  --gin.INITIAL_CHECKPOINT_PATH=\"'{pretrained_path}'\" \\\n",
    "  --gin.TRAIN_STEPS='{train_steps}' \\\n",
    "  --gin.SAVE_PERIOD='{save_period}'\\\n",
    "  --gin.EVAL_PERIOD='{eval_period}'\\\n",
    "  --gin.MIXTURE_OR_TASK_NAME=\"'{task}'\" \\\n",
    "  --gin.TASK_FEATURE_LENGTHS=\"{features_length}\" \\\n",
    "  --gin.BATCH_SIZE='{batch_size}' \\\n",
    "  --task=\"{task}\" \\\n",
    "  --metric=\"{metric}\" \\\n",
    "  --train_file=\"{train_file}\" \\\n",
    "  --predict_file=\"{test_file}\" # or {dev_file}\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "eedb55a3f3d5a08c90a45b02edd9d5201f64a9996f64fdac14a22b56503f46e8"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}