pipeline.py

# -*- coding: utf-8 -*-
"""Pipeline.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1vXr51lAshWxbrrQYNr0JgFsTLnD-T_kL
"""
# Библиотеки
import torch
from tqdm import tqdm
import torchvision
import random
import numpy as np
import locale
def getpreferredencoding(do_setlocale = True):
    return "UTF-8"
locale.getpreferredencoding = getpreferredencoding

from transformers import (
    T5ForConditionalGeneration, T5Tokenizer,
    TrainingArguments, Trainer,
)
from datasets import load_dataset
from rouge import Rouge
from nltk.translate.bleu_score import corpus_bleu
from evaluate import load

# Seed
torch.manual_seed(42)
random.seed(42)
np.random.seed(42)
#torch.use_deterministic_algorithms(True)

# Пути для сохранения:
path = '.'
saved_model_path = path + '/archive'
checkpoint_path = path + '/t5-model-small'
logs_path = checkpoint_path + '/logs'
dataset_path = path + '/dataset'
model_name = "cointegrated/rut5-small"
pretrained_model = saved_model_path + '/model_t5_small_4.pth'
test_path = path + '/texts'

# Устройство ускорителя:
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Функции:
def make_dataset(data, tokenizer, max_length_text=2890, max_length_ref=200):
    '''
    Создать датасет для обучения модели: исходный текс обрабатывается токенизатором, а затем к полученному словарю добавляется метка label
    с токенизированным эталонным рефератом

    Возвращает преобразованный датасет (list)
    '''
    dataset = []
    for inst in tqdm(data):
        txt = tokenizer(inst['text'], add_special_tokens=True, max_length=max_length_text, padding="max_length", truncation=True)
        sum_ = tokenizer(inst['summary'], add_special_tokens=True, max_length=max_length_ref, padding="max_length", truncation=True).input_ids
        txt["labels"] = sum_
        dataset.append(txt)
    return dataset

def save(path, model, optimizer):
    '''
    Сохранить модель и оптимизатор в path
    '''
    torch.save({
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict()
    }, path)

def summarize(text, model, tokenizer, max_length_text=2890, max_length_ref=500):
    '''
    Генерация реферата

    Возвращает сгенерированный моделью реферат (str)
    '''
    inp = tokenizer(text, add_special_tokens=True, max_length=max_length_text, padding="max_length", truncation=True, return_tensors='pt').to(device) 
    return tokenizer.decode(model.generate(input_ids=inp.input_ids, attention_mask=inp.attention_mask, max_length=max_length_ref)[0], skip_special_tokens=True)

def tests_res(data, model, tokenizer, max_length_text=2890):
    '''
    Генерация рефератов для тестирования модели

    Возвращает результат модели на датасете data (list) и эталонные рефераты (list)
    '''
    res, ref = [], []
    for inst in tqdm(data):
        res.append(summarize(inst['text'], model, tokenizer))
        ref.append(inst['summary'])
    return res, ref

def read_dataset(model, tokenizer, max_length_text=2890, max_length_ref=200, n=50):
    '''
    Создать датасет для обучения (для чтения из директории)

    Возвращает преобразованный датасет (list)
    '''
    dataset = []
    for i in range(n):
        with open(dataset_path + f'/{i}.txt', 'r', encoding='utf-8') as f:
            data = f.read()
        data = data.split('\n\n')
        txt = tokenizer(data[1], add_special_tokens=True, max_length=max_length_text, padding="max_length", truncation=True)
        sum_ = tokenizer(data[2], add_special_tokens=True, max_length=max_length_ref, padding="max_length", truncation=True).input_ids
        txt["labels"] = sum_
        dataset.append(txt)
    return dataset

def train_test_model(model, tokenizer, optimizer,
                     train_dataset, val_dataset, test_dataset,
                     num_steps=1, num_epochs=1):
    '''
    Обучение + тестирование + логирование
    '''

    for step in range(1, num_steps+1):
        # Тренер и параметры
        if val_dataset is None:
            training_args = TrainingArguments(
                        output_dir= checkpoint_path,
                        overwrite_output_dir=True,
                        per_device_train_batch_size=2,
                        num_train_epochs=num_epochs,
                        warmup_steps=10,
                        gradient_accumulation_steps=16,
                        save_strategy="epoch",
                        seed=42,
                    )
        
            trainer = Trainer(
                        model=model,
                        args=training_args,
                        train_dataset=train_dataset,
                        tokenizer=tokenizer,
                        optimizers = (optimizer, None)
                    )
        else:
            training_args = TrainingArguments(
                    output_dir= checkpoint_path,
                    overwrite_output_dir=True,
                    per_device_train_batch_size=2,
                    per_device_eval_batch_size=2,
                    num_train_epochs=num_epochs,
                    warmup_steps=10,
                    gradient_accumulation_steps=16,
                    evaluation_strategy="no",
                    save_strategy="epoch",
                    seed=42,
                )

            trainer = Trainer(
                    model=model,
                    args=training_args,
                    train_dataset=train_dataset,
                    eval_dataset=val_dataset,
                    tokenizer=tokenizer,
                    optimizers = (optimizer, None)
                )
        
        # Обучение модели
        model.train()
        logs_train = trainer.train()
        logs_eval = None
        if val_dataset is not None:
            model.eval()
            logs_eval = trainer.evaluate()
    
        print('Saving model')
        save(saved_model_path + f'/model_t5_small_6_{step}.pth', model, optimizer)
    
        print('Saving loss')
        with open(logs_path + f'/loss_dictionary.txt','a') as f:
            f.write(f'Step_{step}\nTRAIN LOG:\n{logs_train}\nEVAL LOG:\n{logs_eval}\n\n')

        # Тестирование модели
        model.eval()
        print('Testing')
        model_results, refs = tests_res(test_dataset, model, tokenizer) # результаты работы модели

        # Оценка ROUGE
        print('Done! Counting Rouge')
        scores = rouge.get_scores(model_results, refs, avg=True)
        print(scores)

        # Оценка BLEU
        print('Done! Counting BLEU')
        blue = corpus_bleu([[r.split(" ")] for r in refs], [hyp.split(" ") for hyp in model_results])
        print(blue)

        # Оценка METEOR
        print('Done! Counting METEOR')
        results_m = meteor.compute(predictions=model_results, references=refs)
        print(results_m)

        # Оценка BertScore
        print('Done! Counting BertScore')
        results_b = bertscore.compute(predictions=model_results, references=refs, lang="ru")
        results_b = {k: np.mean(v) for k, v in list(results_b.items())[:-1]}
        print(results_b)

        print('Saving scores')
        with open(logs_path + f'/metrics.txt', 'a') as f:
            f.write(f'STEP: {step}\n')
            f.write(f'ROUGE: {scores}\n')
            f.write(f'BLEU: {blue}\n')
            f.write(f'METEOR: {results_m}\n')
            f.write(f'BertScore: {results_b}\n\n')

def for_human_eval(model, n=64):
    '''
    Печать + сохранение результатов в файл
    '''    
    model.eval()
    res = ''
    for i in range(n):
        with open(test_path + f'/text_{i}', 'r', encoding='cp1251') as f:
            t = f.read()
        new = f'{i}) {summarize(t)}'
        print(new)
        print('-----------')
        res += new + '\n-----------\n'
    with open(test_path + f'/results.txt', 'w') as f:
        f.write(res)


if __name__ == '__main__':
    
    # Загрузка модели
    tokenizer = T5Tokenizer.from_pretrained(model_name)
    model_t5 = T5ForConditionalGeneration.from_pretrained(model_name)
    optimizer = torch.optim.AdamW(model_t5.parameters(),lr=1e-4)
    
    checkpoint = torch.load(pretrained_model, map_location='cpu')
    model_t5.load_state_dict(checkpoint['model_state_dict'])
    model_t5.to(device)
    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
    
    # Загрузка датасета
    dataset_train = load_dataset('IlyaGusev/gazeta', revision="v2.0")["train"]
    dataset_val = load_dataset('IlyaGusev/gazeta', revision="v2.0")["validation"]
    dataset_test = load_dataset('IlyaGusev/gazeta', revision="v2.0")["test"]
    
    # Создание обучающей и валидационной выборок
    print('Making train dataset')
    dataset_train = make_dataset(dataset_train, tokenizer)
    print('Making val dataset')
    dataset_val = make_dataset(dataset_val, tokenizer)
    print('Done!')
    
    # Загрузка метрик
    rouge = Rouge()
    meteor = load('meteor')
    bertscore = load("bertscore")
    
    ### Дообучение на датасете Gazeta:
    
    num_steps = 1
    
    train_test_model(model_t5, tokenizer, optimizer,
                     dataset_train, dataset_val, dataset_test,
                     num_steps=num_steps)
    
    ## Тестирование на своей выборке текстов:
    
    for_human_eval(model_t5)
    
    ### Дообучение на своем датасете:
    
    # Загрузка датасета и создание выборок
    train_set = read_dataset(model_t5, tokenizer)
    
    test_set = []
    for i in range(50, 70):
        with open(dataset_path + f'/{i}.txt', 'r', encoding='utf-8') as f:
            data = f.read()
            data = data.split('\n\n')
            test_set.append({'text': data[1], 'summary': data[2]})
    
    #Обучение + тестирование:    
    num_steps = 1
    
    train_test_model(model_t5, tokenizer, optimizer,
                     train_set, None, test_set,
                     num_steps=num_steps)
    
    ### Тестирование на своей выборке текстов:
    
    for_human_eval(model_t5)