text_expansion.py

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

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1-iF4CjDh_LHQrFlhgQjUEWsk5Cz6ObiF

# AI CopyWriter

Прототип перефразирования и/или распространения текста на основе ruGPT3Large.

### Установим необходимые зависимости:
"""

# !pip3 install urllib3==1.25.4

# !pip3 install transformers==2.8.0

# !pip install pymorphy2==0.8

# !pip install spacy==2.1.9

# !python3 -m spacy download xx_ent_wiki_sm

# !pip install deeppavlov

# !pip3 uninstall tensorflow --yes

# !pip3 install tensorflow-gpu==1.15.2

# !python3 -m  deeppavlov install syntax_ru_syntagrus_bert

# ! pip3 install jsonlines

"""### Now several cells with the code."""

import jsonlines
import json
import os
import random
import re
import string  
from copy import deepcopy
from pathlib import Path

import numpy as np
import pandas as pd
import pymorphy2
import spacy
import xx_ent_wiki_sm
from deeppavlov import build_model, configs

morph = pymorphy2.MorphAnalyzer()

nlp = xx_ent_wiki_sm.load()
nlp.add_pipe(nlp.create_pipe('sentencizer'), first=True)

syntax_model = build_model(configs.syntax.syntax_ru_syntagrus_bert, download=True)

from transformers import GPT2LMHeadModel,GPT2Tokenizer
import torch

tokenizer = GPT2Tokenizer.from_pretrained("sberbank-ai/rugpt3large_based_on_gpt2")

model = GPT2LMHeadModel.from_pretrained("sberbank-ai/rugpt3large_based_on_gpt2")
if torch.cuda.is_available():
  device = "cuda"
else:
  device = "cpu"
model.to(device)
print("Success!")

russian_restricted_pronouns = "я мной меня мною мне мы нас нам нами ты тебя тебе тобою тобой вы вас вам вами".split()
extra_marks = re.compile(r"&[a-zA-Z0-9;]+")
expanding_startings = [
    "В то же время", 
    "Это так", 
    "Действительно,", 
    "Потому что", 
    "А главное,", 
    "При этом",
    "В связи с этим",
    "Кроме того,",
    "Интересно, что",
    "Также",
    ]


def join_words_in_or_pattern(words):
    return "(" + "|".join([r'\b%s\b' % word for word in words]) + ")"


RU_PRONOUNS = re.compile(join_words_in_or_pattern(russian_restricted_pronouns), re.IGNORECASE)


def generate_rugpt3large(prompt_text, return_only_predicted=False,
                         till_new_string=True, generate_only_one_sent=True, 
                         length=100, temperature=1.0, k=10, p=0.9, repetition_penalty=1.0,
                         num_return_sequences=1, stop_token="</s>"):
    global device
    total_sequence = ""
    if num_return_sequences > 1:
        all_sequences = []
    
    encoded_prompt = tokenizer.encode(
        prompt_text, add_special_tokens=False, return_tensors="pt")
    encoded_prompt = encoded_prompt.to(device)

    output_sequences = model.generate(
                input_ids=encoded_prompt,
                max_length=length + len(encoded_prompt[0]),
                temperature=temperature,
                top_k=k,
                top_p=p,
                repetition_penalty=repetition_penalty,
                do_sample=True,
                num_return_sequences=num_return_sequences,
            )
    
    if len(output_sequences.shape) > 2:
        output_sequences.squeeze_()

    for generated_sequence_idx, generated_sequence in enumerate(output_sequences):
        generated_sequence = generated_sequence.tolist()
        
        text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
        text = text[: text.find(stop_token) if stop_token else None]

        total_sequence = (text[len(tokenizer.decode(
            encoded_prompt[0], clean_up_tokenization_spaces=True)) :])
        total_sequence = re.sub(r"[\"'«»]", "", total_sequence)
        total_sequence = total_sequence.strip()
        
        if till_new_string:
            total_sequence = total_sequence[:total_sequence.find("\n")].strip()
            
        if generate_only_one_sent:
            total_sequence = " ".join([sent.text for sent in nlp(total_sequence).sents][:1])
            
        if len(total_sequence) > 1 and prompt_text[-1] in [".", "!", "?", "…"]:
            total_sequence = total_sequence[0].upper() + total_sequence[1:]
            
        if not return_only_predicted:
            if len(total_sequence) > 1 and total_sequence[0] in [".", "!", "?", "…", ","]:
                total_sequence = prompt_text + total_sequence
            else:
                total_sequence = prompt_text + " " + total_sequence
        else:
            if len(total_sequence) > 1 and total_sequence[0] in [".", "!", "?", "…", ","]:
                total_sequence = total_sequence
            else:
                total_sequence = " " + total_sequence
        
        if num_return_sequences > 1 and len(total_sequence) > 0:
            all_sequences.append(total_sequence)
            
    if num_return_sequences > 1:
        return all_sequences
    else:
        return total_sequence
    

def get_nsubjects(text):
    nsubjects = []
    for parse in syntax_model([text]):
        for row in parse.split("\n"):
            if "nsubj" in row:
                nsubjects.append(row.split("\t")[1])

    return nsubjects


def is_satisfying(sent):
    doc = nlp(sent)
    if len(doc.ents) > 0:
        return False
    
    ntokens = len(sent.split())
    if 15 <= ntokens or ntokens < 5:
        return False
    if re.search(RU_PRONOUNS, sent):
        return False
    
    return True


def generate_paraphrase(text):
    predictions = generate_rugpt3large(f"{text} Перефразирую:", 
                                       return_only_predicted=True, 
                                       num_return_sequences=10)
    predictions = [p.strip() for p in predictions]
    predictions = [p for p in predictions if len(p) > 1]
    original_tokens = set(text.split())
    nsubjects = set(get_nsubjects(text))
    
    total_n_tokens = []
    total_same_tokens = []
    total_same_nsubjects = []
    
    for pred_par in predictions:
        pred_tokens = pred_par.split()
        
        len_pred_tokens = len(pred_tokens)
        len_same_tokens = len(original_tokens.intersection(set(pred_tokens)))
        len_same_nsubjs = len(nsubjects.intersection(set(pred_tokens)))
        
        total_n_tokens.append(len_pred_tokens)
        total_same_tokens.append(len_same_tokens)
        total_same_nsubjects.append(len_same_nsubjs)
    
    targets = [st * sn / tt for st, sn, tt in 
               zip(total_same_tokens, total_same_nsubjects, total_n_tokens)]
    # if all([el <= 1 for el in total_same_nsubjects]):
    if sum(total_same_nsubjects) == 0:
        targets = [st / tt for st, sn, tt in 
                   zip(total_same_tokens, total_same_nsubjects, total_n_tokens)]
    targets = [t if predictions[i] != text else 0 for i, t in enumerate(targets)]
    best_pred_id = np.argmax(targets)
    
    new_sent_text = predictions[best_pred_id]
    if len(new_sent_text) > 1:
        if new_sent_text[-1] not in [".", "!", "?", "…"]:
            new_sent_text = new_sent_text + "."
        new_sent_text = new_sent_text[0].upper() + new_sent_text[1:]
    else:
        new_sent_text = text
    return new_sent_text

    
def expand_text(text, paraphrase=False, max_history_sents=5):
    expanded = deepcopy(text)
    
    doc = nlp(text)
    expanded_sents = []
    sents_texts = [sent.text for sent in doc.sents if len(sent.text) > 1]
    for sent_id, sent in enumerate(doc.sents):
        if len(sent.text) <= 1:
            continue
        expanded_sents.append(sent.text)
        
        expand_sent = ""
        if sent_id < len(sents_texts) - 1 and {"CONJ"} in morph.parse(sents_texts[sent_id + 1].split()[0])[0].tag:
            pass
        else:
            expanding_start = random.choice(expanding_startings)
            context = " ".join(expanded_sents[-max_history_sents:])
            generated = generate_rugpt3large(f"{context} {expanding_start}",
                                             return_only_predicted=True, num_return_sequences=5)
            
            for gen in generated:
                if is_satisfying(gen):
                    expand_sent = f"{expanding_start}{gen}".strip()
                    if expand_sent[-1] not in [".", "!", "?", "…"]:
                        expand_sent = expand_sent + "."
                    break
        
        new_subtext = deepcopy(sent.text)
        if paraphrase:
            new_sent_text = generate_paraphrase(sent.text).strip()
            if len(new_sent_text) > 1:
                if new_sent_text[-1] not in [".", "!", "?", "…"]:
                    new_sent_text = new_sent_text + "."
                new_sent_text = new_sent_text[0].upper() + new_sent_text[1:]
            else:
                new_sent_text = deepcopy(sent.text)
                
            if len(expand_sent) > 0:
                new_subtext = f"{new_sent_text} {expand_sent}"
                expanded_sents.append(f"{expand_sent}")
            else:
                new_subtext = f"{new_sent_text}"
        else:
            if len(expand_sent) > 0:
                new_subtext = f"{sent.text} {expand_sent}"
                expanded_sents.append(f"{expand_sent}")
                
        expanded = expanded.replace(sent.text, new_subtext)
        
                
    expanded = expanded.strip()
    if abs(len(expanded) - len(text)) < 10:
        expanded = expand_text(text, max_history_sents=5)
    return expanded

def paraphrase_and_expand_text(text, paraphrase=False, expand=False, 
                               max_history_sents=5):
    text = text.strip()
    if expand:
        paraphrased_expanded = expand_text(text, 
                                           paraphrase=paraphrase, 
                                           max_history_sents=max_history_sents)
    elif paraphrase:
        paraphrased_expanded = ""
        doc = nlp(text)
        for sent_id, sent in enumerate(doc.sents):
            if len(sent.text) <= 1:
                continue

            if paraphrase:
                par_sent = generate_paraphrase(sent.text)
                paraphrased_expanded += f" {par_sent}"
    paraphrased_expanded = re.sub(extra_marks, "", paraphrased_expanded)
    return paraphrased_expanded.strip()

"""# DEMO Usage

Итак, все зависимости установлены, функции определены. Самое время просто воспользоваться DEMO. 
Пользователь демонстрации может изменять список текстов, и главное параметры `paraphrase` и `expand`. Параметр `paraphrase` отвечает за то, будет ли исходный текст перефразирован, а параметр `expand` показывает следует ли модели дополнить текст. Эти режимы могут быть включены как по отдельности, так и вместе.
"""

for text in [
             """На мой взгляд, лето – прекраснейшее время года, поскольку очень тепло, вокруг все красиво, и у нас имеется отличная возможность отдохнуть. 
             А еще лето – это время каникул и отпусков. 
             Мы можем отправиться куда угодно, замечательно провести время у водоема, в тени изумрудной зелени, плескаться в теплой воде. 
             А можно отправиться в горы, носиться на велосипеде или играть во всевозможные игры.""",
             """Лето – это всегда особенное время года и каждый человек старается сделать его запоминающимся. 
             Тут пригодится умение замечать простые мелочи, впитывать в себя солнечные и радостные моменты каждого дня. 
             Даже если на улице идет дождь, поездка на море отменяется, а любимое мороженное в магазине закончилось. 
             Ведь это всего лишь отдельные дни и в целом лето это калейдоскоп ярких переживаний, который мы держим в своих руках.""",
             """Тысячи капель прохладной воды падают на лицо и стекают по волосам и телу, солнце нежно щекотит кожу, а ты вдыхаешь запах сочной листвы. 
             И вы не просто мажете тело обёртыванием или наносите на лицо сыворотку, вы совершаете ритуал любви к себе, своей коже и своему телу. 
             Это волшебное время, когда девушка находиться наедине с собой. Она вся в своих мыслях, она приводит не только себя в порядок, но и свой разум. 
             Восстанавливается после тяжелого дня или наоборот, только готовится к нему, чтобы быть бодрой и энергичной."""]:

    rewritten_text = paraphrase_and_expand_text(text, paraphrase=True, expand=True)
    print(f"Original Text: {text}\n\nRewritten Text: {rewritten_text}\n\n\n")