import streamlit as st # !pip install -q transformers import numpy as np # import pandas as pd import re # import random import torch # from tqdm.notebook import tqdm import transformers # from torch.optim import AdamW import textwrap # Загружаем токенайзер модели from transformers import GPT2Tokenizer tokenizer = GPT2Tokenizer.from_pretrained('sberbank-ai/rugpt3small_based_on_gpt2') # import re with open('anekdoty.txt', encoding='utf8') as f: text = f.read() text = re.sub('\n{2,}', '\n', text) print(text[:1000]) # токенизируем текст tokens = tokenizer.encode(text, add_special_tokens=True) tokens = np.array(tokens) print(len(tokens)) tokens[:10] # разбиваем на train и test l = len(tokens)//15 train = [] test = [] for i in range(15): if i%5 > 0: train.extend(tokens[i*l: (i+1)*l]) else: test.extend(tokens[i*l: (i+1)*l]) train = np.array(train) test = np.array(test) print(len(tokens), len(train), len(test)) from transformers import GPT2LMHeadModel # Эту модель просто подгружаем и не будем дообучать model_init = GPT2LMHeadModel.from_pretrained( 'sberbank-ai/rugpt3small_based_on_gpt2', output_attentions = False, output_hidden_states = False, ) # Эту модель подгрузим и далее обучим model = GPT2LMHeadModel.from_pretrained( 'sberbank-ai/rugpt3small_based_on_gpt2', output_attentions = False, output_hidden_states = False, ) model.to(device); model_init.to(device); batch_size = 8 max_len = 256 epochs = 5 n_train = len(train)//(batch_size*max_len) n_test = len(test)//(batch_size*max_len) print(n_train, n_test) # устанавливаем оптимизатор optimizer = AdamW(model.parameters(), lr = 1e-5, eps = 1e-8) # трансформеры с трудом обучаются, для них нужны разные способы повышения # эффективности градиентного спуска total_steps = n_train * epochs scheduler = transformers.get_linear_schedule_with_warmup(optimizer, num_warmup_steps = 0, num_training_steps = total_steps) # зададим точность, хотя ориентироваться будем на качество генерации def accuracy(y_true, logits): return torch.mean((y_true[1:] == torch.argmax(logits, dim=2)[:-1]).float()).detach().cpu().numpy() # готовим тензоры для обучения размера [batch_size, max_len] def prep_tensors(x, i, batch_size=batch_size, max_len=max_len): batch_ids = x[i*batch_size*max_len: (i+1)*batch_size*max_len] batch_ids = batch_ids.reshape(batch_size, max_len) batch_ids = torch.tensor(batch_ids).to(device) return batch_ids # обучающий цикл for epoch in range(1, epochs+1): print(f'epoch {epoch}/{epochs} : training') train_loss = [] train_acc = [] model.train() pbar = range(n_train) # pbar = tqdm(range(n_train)) for i in pbar: batch_ids = prep_tensors(train, i) model.zero_grad() loss, logits, _ = model(batch_ids, token_type_ids=None, labels=batch_ids ).values() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() scheduler.step() train_loss.append(loss.item()) train_acc.append(accuracy(batch_ids, logits)) print(f'acc {np.mean(train_acc):.4f} loss {np.mean(train_loss):.4f}') # pbar.set_description(f'acc {np.mean(train_acc):.4f} loss {np.mean(train_loss):.4f}', refresh=True) print(f'epoch {epoch}/{epochs} : validation') model.eval() val_acc = [] val_loss = [] pbar = range(n_test) # pbar = tqdm(range(n_test)) for i in pbar: batch_ids = prep_tensors(test, i) with torch.no_grad(): loss, logits, _ = model(batch_ids, token_type_ids=None, labels=batch_ids ).values() val_loss.append(loss.item()) val_acc.append(accuracy(batch_ids, logits)) print(f'acc {np.mean(val_acc):.4f} loss {np.mean(val_loss):.4f}') # pbar.set_description(f'acc {np.mean(val_acc):.4f} loss {np.mean(val_loss):.4f}', refresh=True) # Применим модель, которую мы не дообучали: просто для понимания разницы между дообученной на собственных данных моделью и предобученной. # https://huggingface.co/transformers/main_classes/model.html#transformers.generation_utils.GenerationMixin.generate # модель без дообучения # prompt – строка, которую модель примет на вход и продолжит prompt = 'Мужик спрашивает официанта' # токенизируем строку prompt = tokenizer.encode(prompt, return_tensors='pt').to(device) # out будет содержать результаты генерации в виде списка out = model_init.generate( # входная строка input_ids=prompt, # максимальная длина генерируемой последовательности max_length=250, # num_beams num_beams=5, # применяем сэмплирование do_sample=True, # применяем температуру temperature=55., # топ слов по вероятности top_k=50, # топ слов по суммарной вероятности top_p=0.6, # сколько (постараться) не повторять n_gram подряд no_repeat_ngram_size=3, # сколько вернуть генераций num_return_sequences=7, ).cpu().numpy() # out содержит результаты # декодируем и печатаем for out_ in out: print(tokenizer.decode(out_)) # дообученная модель with torch.inference_mode(): prompt = 'Мужик спрашивает официанта' prompt = tokenizer.encode(prompt, return_tensors='pt').to(device) out = model.generate( input_ids=prompt, max_length=150, num_beams=1, do_sample=True, temperature=1., top_k=5, top_p=0.6, no_repeat_ngram_size=2, num_return_sequences=7, ).cpu().numpy() for out_ in out: print(textwrap.fill(tokenizer.decode(out_), 100), end='\n------------------\n') # Сохраняем веса обученной модели torch.save(model.state_dict(), 'model.pt') # Задаем класс модели (уже в streamlit/tg_bot) model_finetuned = GPT2LMHeadModel.from_pretrained( 'sberbank-ai/rugpt3small_based_on_gpt2', output_attentions = False, output_hidden_states = False, ) # Вешаем сохраненные веса на нашу модель model = model_finetuned.load_state_dict(torch.load('model.pt')) # ->