tensorflow_official_nlp_transformer使用示例

github链接

https://github.com/tensorflow/models/tree/master/official/nlp/transformer

安装

pip install tf-models-nightly （readme有介绍，也可以源码安装）

使用说明

用自有数据训练预测，没有严格按照官方教程走流程，有些步骤解耦出来处理的

1. 导入必要的包

import os
import logging
import numpy as np
import tensorflow as tf
from official.nlp.transformer import data_download
from official.nlp.transformer.utils import tokenizer
from official.nlp.transformer import misc, transformer_main, transformer, translate
from official.utils.flags import core as flags_core

2. 数据预处理

raw_dir = "./data_dir/tf2_nlp/"
data_dir = "./data_dir/tf2_nlp/"
vocab_file = './data_dir/tf2_nlp/vocab.enzh'

TRAIN_DATA = [{
   "url": "", "input": "raw_en.txt", "target": "raw_zh.txt"}]
EVAL_DATA = [{
   "url": "", "input": "eval_en.txt", "target": "eval_zh.txt"}]

# 获取原始数据地址，可省略，直接写文件地址
logging.info("Step 1: Downloading data from source")
train_files = data_download.get_raw_files(raw_dir, TRAIN_DATA)
eval_files = data_download.get_raw_files(raw_dir, EVAL_DATA)

# 构建分词器
logging.info("Step2: Creating subtokenizer and building vocabulary")
train_files_flat = train_files["inputs"] + train_files["targets"]
subtokenizer = tokenizer.Subtokenizer.init_from_files(
      vocab_file, train_files_flat