import torch
import numpy as np
from torch.nn import functional as F
from transformers import GPT2Tokenizer, GPT2LMHeadModel
def top_filtering(logits, top_k):
"""
Фильтрация top-k, в фильтрации top-p в этой задаче особо смысла нет
код с top-p: https://github.com/ictnlp/DSTC8-AVSD/blob/master/generate.py
"""
assert logits.dim() == 1
top_k = min(top_k, logits.size(-1))
if top_k > 0:
indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
logits[indices_to_remove] = -float('Inf')
return logits
class TextEmbdGenerator:
def __init__(self, name_or_path, sbert, device = None):
"""
Инициализация генератора текста с моделью и токенизатором.
name_or_path: путь до модели токенизатора или ее имя для загрузки из Hugging Face.
sbert: модель для ранжирования (такая же что и создает эбеддинги)
"""
self.device = device
if self.device == None:
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.tokenizer = GPT2Tokenizer.from_pretrained(name_or_path)
self.model = GPT2LMHeadModel.from_pretrained(name_or_path).to(self.device)
self.sbert = sbert
def __get_embds(self, embds, sqr_embds):
'''Работает с матрицей эмбеддингов'''
list_emb = []
sq = embds if sqr_embds == None else sqr_embds
for i, embd in enumerate(embds):
vector = np.concatenate([embd,sq[i]**2])
list_emb.append(list(vector))
return torch.FloatTensor(list_emb).to(self.device)
def __det_mean(self, embds):
'''Получение среднего'''
m = np.zeros((384))
for embd in embds:
m += embd
m /= len(embds)
return m
def generate_embedding(self, embds, prompt = '', sqr_embds = None, temperature=0.26, top_k=4, max_len=100):
"""
Генерация текста на основе начального эмбеддинга и заданного начального текста.
"""
current_output_ids = self.tokenizer.encode(prompt)
embedding = self.__get_embds(embds, sqr_embds)
word_tokens = self.model.base_model.wte
while len(current_output_ids) < max_len:
with torch.no_grad():
outp_ids_tensor = torch.LongTensor(current_output_ids).to(self.device)
token_embeddings = word_tokens(outp_ids_tensor)
input_vectors = torch.vstack((embedding, token_embeddings)).unsqueeze(dim=0)
output_model = self.model(inputs_embeds=input_vectors)
logits = output_model.logits
if isinstance(logits, tuple):
logits = logits[0]
logits = logits[0, -1, :]
logits /= temperature
logits = top_filtering(logits, top_k)
probs = F.softmax(logits, dim=-1)
prev = torch.multinomial(probs, 1)
if prev.item() == self.tokenizer.eos_token_id:
break
current_output_ids.append(prev.item())
output_text = self.tokenizer.decode(current_output_ids)
return output_text.split('\n')[0]
def cosine_similarity(self, x, y):
"""Вычисление косинусного сходства."""
return np.dot(x, y) / (np.linalg.norm(x) * np.linalg.norm(y))
def generate_with_ranker(self, embds, prompt = '', sqr_embds = None, seq=10, temperature=0.6, top_k=10, max_len=100):
"""Генерация и ранжирование текста. Поумолчанию создаются 10 текстов"""
sequences = [self.generate_embedding(embds, prompt, sqr_embds, temperature, top_k, max_len) for _ in range(seq)]
sequences = list(set(sequences))
embd = self.__det_mean(embds)
embeddings = self.sbert.encode(sequences)
similarities = [self.cosine_similarity(embd, emb) for emb in embeddings]
best_index = np.argmax(similarities)
return sequences[best_index]