T5 Russian Spell

🚀 t5-russian-spell

这是一个用于纠正从音频识别文本的模型。可以将我的音频识别模型（https://huggingface.co/UrukHan/wav2vec2-russian ）的识别结果输入到这个模型中。我已经在随机选取的YouTube视频上进行了测试。

🚀 快速开始

运行示例代码

你可以在Colab中运行以下示例代码来查看模型的运行结果：点击此处查看Colab代码

# 安装transformers库
!pip install transformers

# 导入库
from transformers import AutoModelForSeq2SeqLM, T5TokenizerFast

# 指定从模型中心选择的模型名称
MODEL_NAME = 'UrukHan/t5-russian-spell'
MAX_INPUT = 256

# 加载模型和分词器
tokenizer = T5TokenizerFast.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

# 输入数据（可以是短语数组或文本）
input_sequences = ['сеглдыя хорош ден', 'когд а вы прдет к нам в госи']   # 也可以使用单个短语:  input_sequences = 'сеглдыя хорош ден'

task_prefix = "Spell correct: "                 # 数据分词
if type(input_sequences) != list: input_sequences = [input_sequences]
encoded = tokenizer(
  [task_prefix + sequence for sequence in input_sequences],
  padding="longest",
  max_length=MAX_INPUT,
  truncation=True,
  return_tensors="pt",
)

predicts = model.generate(encoded)    # 预测

tokenizer.batch_decode(predicts, skip_special_tokens=True)  # 解码数据

运行训练和保存模型

你可以使用以下配置好的笔记本在Hugging Face Hub上运行训练并将模型保存到自己的仓库中：点击此处查看Colab代码

# 安装库
!pip install datasets
!apt install git-lfs
!pip install transformers
!pip install sentencepiece 
!pip install rouge_score

# 导入库
import numpy as np
from datasets import Dataset
import tensorflow as 
import nltk
from transformers import T5TokenizerFast, Seq2SeqTrainingArguments, Seq2SeqTrainer, AutoModelForSeq2SeqLM, DataCollatorForSeq2Seq
import torch
from transformers.optimization import Adafactor, AdafactorSchedule
from datasets import load_dataset, load_metric

# 加载参数
raw_datasets = load_dataset("xsum")
metric = load_metric("rouge")
nltk.download('punkt')

# 输入你的Hugging Face Hub密钥
from huggingface_hub import notebook_login
notebook_login()

# 定义参数
REPO = "t5-russian-spell"  # 输入仓库名称
MODEL_NAME = "UrukHan/t5-russian-spell" # 输入从模型中心选择的模型名称
MAX_INPUT = 256  # 输入最大输入数据长度（以标记为单位）（输入短语的单词长度（可以将半个单词视为一个标记））
MAX_OUTPUT  = 256 # 输入最大预测长度（以标记为单位）（对于摘要任务或其他输出较短的任务，可以减小此值）
BATCH_SIZE = 8 
DATASET = 'UrukHan/t5-russian-spell_I'   # 输入数据集名称

# 加载数据集，使用其他类型的数据将在下面描述
data = load_dataset(DATASET)

# 加载模型和分词器
tokenizer = T5TokenizerFast.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

model.config.max_length = MAX_OUTPUT  # 默认值为20，因此在所有模型中，输出序列都会被截断
# 第一次保存到仓库后可以注释掉，这不是必需的
tokenizer.push_to_hub(repo_name) 

train = data['train']
test = data['test'].train_test_split(0.02)['test']  # 减小测试集，以免在各轮训练之间等待很长时间来计算误差

data_collator = DataCollatorForSeq2Seq(tokenizer, model=model) #return_tensors="tf"

def compute_metrics(eval_pred):
  predictions, labels = eval_pred
  decoded_preds = tokenizer.batch_decode(predictions, skip_special_tokens=True)
  # 替换标签中的-100，因为我们无法对其进行解码
  labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
  decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
  
  # Rouge要求每个句子后有一个换行符
  decoded_preds = ["\n".join(nltk.sent_tokenize(pred.strip())) for pred in decoded_preds]
  decoded_labels = ["\n".join(nltk.sent_tokenize(label.strip())) for label in decoded_labels]
  
  result = metric.compute(predictions=decoded_preds, references=decoded_labels, use_stemmer=True)
  # 提取一些结果
  result = {key: value.mid.fmeasure * 100 for key, value in result.items()}
  
  # 添加平均生成长度
  prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in predictions]
  result["gen_len"] = np.mean(prediction_lens)
  
  return {k: round(v, 4) for k, v in result.items()}
  
training_args = Seq2SeqTrainingArguments(
  output_dir = REPO,
  #overwrite_output_dir=True,
  evaluation_strategy='steps',
  #learning_rate=2e-5,
  eval_steps=5000,
  save_steps=5000,
  num_train_epochs=1,
  predict_with_generate=True,
  per_device_train_batch_size=BATCH_SIZE,
  per_device_eval_batch_size=BATCH_SIZE,
  fp16=True,
  save_total_limit=2,
  #generation_max_length=256,
  #generation_num_beams=4,
  weight_decay=0.005,
  #logging_dir='logs',
  push_to_hub=True,
)

# 手动选择优化器。T5在原始架构中使用Adafactor优化器
optimizer = Adafactor(
    model.parameters(),
    lr=1e-5,
    eps=(1e-30, 1e-3),
    clip_threshold=1.0,
    decay_rate=-0.8,
    beta1=None,
    weight_decay=0.0,
    relative_step=False,
    scale_parameter=False,
    warmup_init=False,
)
lr_scheduler = AdafactorSchedule(optimizer)

trainer = Seq2SeqTrainer(
  model=model,
  args=training_args,
  train_dataset = train,
  eval_dataset = test,
  optimizers = (optimizer, lr_scheduler),
  tokenizer = tokenizer,
  compute_metrics=compute_metrics
)

trainer.train()

trainer.push_to_hub()

数据转换示例

以下是将数组转换为适合此网络格式的示例代码：

input_data = ['удач почти отнее отвернулась', 'в хааоде проведения чемпиониавта мира дветысячивосемнандцтая лгодаа']
output_data = ['Удача почти от нее отвернулась', 'в ходе проведения чемпионата мира две тысячи восемнадцатого года']

# 对输入数据进行分词
task_prefix = "Spell correct: "
input_sequences = input_data 
encoding = tokenizer(
  [task_prefix + sequence for sequence in input_sequences],
  padding="longest",
  max_length=MAX_INPUT,
  truncation=True,
  return_tensors="pt",
)
input_ids, attention_mask = encoding.input_ids, encoding.attention_mask

# 对输出数据进行分词
target_encoding = tokenizer(output_data, padding="longest", max_length=MAX_OUTPUT, truncation=True)
labels = target_encoding.input_ids
# 将标签中的填充标记ID替换为-100
labels = torch.tensor(labels)
labels[labels == tokenizer.pad_token_id] = -100

# 将我们的数据转换为数据集格式
data = Dataset.from_pandas(pd.DataFrame({'input_ids': list(np.array(input_ids)), 'attention_mask': list(np.array(attention_mask)), 'labels': list(np.array(labels))}))
data = data.train_test_split(0.02)
# 然后将其输入到我们的训练器中：   train_dataset = data['train'],  eval_dataset = data['test']

✨ 主要特性

该模型能够有效纠正音频识别文本中的拼写错误，以下是一个对比示例：

📦 安装指南

在运行代码前，你需要安装以下依赖库：

!pip install datasets
!apt install git-lfs
!pip install transformers
!pip install sentencepiece 
!pip install rouge_score

📚 详细文档

训练数据集

以下是用于训练该模型的数据集：

• UrukHan/t5-russian-spell_I: 数据集链接
• UrukHan/t5-russian-spell_II: 数据集链接
• UrukHan/t5-russian-spell_III: 数据集链接