Views
No views yet
1import torch
2import numpy as np
3from torch.nn import functional as F
4from transformers import GPT2Tokenizer, GPT2LMHeadModel
5
6def top_filtering(logits, top_k):
7 """
8 Фильтрация top-k, в фильтрации top-p в этой задаче особо смысла нет
9 код с top-p: https://github.com/ictnlp/DSTC8-AVSD/blob/master/generate.py
10 """
11 assert logits.dim() == 1
12 top_k = min(top_k, logits.size(-1))
13 if top_k > 0:
14 indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
15 logits[indices_to_remove] = -float('Inf')
16
17 return logits
18
19
20class TextEmbdGenerator:
21 def __init__(self, name_or_path, sbert, device = None):
22 """
23 Инициализация генератора текста с моделью и токенизатором.
24 name_or_path: путь до модели токенизатора или ее имя для загрузки из Hugging Face.
25 sbert: модель для ранжирования (такая же что и создает эбеддинги)
26 """
27 self.device = device
28
29 if self.device == None:
30 self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
31
32
33 self.tokenizer = GPT2Tokenizer.from_pretrained(name_or_path)
34 self.model = GPT2LMHeadModel.from_pretrained(name_or_path).to(self.device)
35 self.sbert = sbert
36
37
38 def __get_embds(self, embds, sqr_embds):
39 '''Работает с матрицей эмбеддингов'''
40 list_emb = []
41 sq = embds if sqr_embds == None else sqr_embds
42
43 for i, embd in enumerate(embds):
44 vector = np.concatenate([embd,sq[i]**2])
45 list_emb.append(list(vector))
46
47 return torch.FloatTensor(list_emb).to(self.device)
48
49 def __det_mean(self, embds):
50 '''Получение среднего'''
51 m = np.zeros((384))
52
53 for embd in embds:
54 m += embd
55
56 m /= len(embds)
57 return m
58
59 def generate_embedding(self, embds, prompt = '', sqr_embds = None, temperature=0.26, top_k=4, max_len=100):
60 """
61 Генерация текста на основе начального эмбеддинга и заданного начального текста.
62 """
63
64 current_output_ids = self.tokenizer.encode(prompt) # Промпт
65 embedding = self.__get_embds(embds, sqr_embds) # Матрица входа
66 word_tokens = self.model.base_model.wte # Словарь токенов
67
68 while len(current_output_ids) < max_len:
69 with torch.no_grad():
70 outp_ids_tensor = torch.LongTensor(current_output_ids).to(self.device) # Выходы
71 token_embeddings = word_tokens(outp_ids_tensor) # эмбеддинги
72 input_vectors = torch.vstack((embedding, token_embeddings)).unsqueeze(dim=0)
73 output_model = self.model(inputs_embeds=input_vectors)
74
75 logits = output_model.logits
76 if isinstance(logits, tuple):
77 logits = logits[0]
78 logits = logits[0, -1, :]
79 logits /= temperature
80 logits = top_filtering(logits, top_k)
81 probs = F.softmax(logits, dim=-1)
82
83 prev = torch.multinomial(probs, 1)
84 if prev.item() == self.tokenizer.eos_token_id:
85 break
86 current_output_ids.append(prev.item())
87
88 output_text = self.tokenizer.decode(current_output_ids)
89 return output_text.split('\n')[0]
90
91
92 def cosine_similarity(self, x, y):
93 """Вычисление косинусного сходства."""
94 return np.dot(x, y) / (np.linalg.norm(x) * np.linalg.norm(y))
95
96 def generate_with_ranker(self, embds, prompt = '', sqr_embds = None, seq=10, temperature=0.6, top_k=10, max_len=100):
97 """Генерация и ранжирование текста. Поумолчанию создаются 10 текстов"""
98 sequences = [self.generate_embedding(embds, prompt, sqr_embds, temperature, top_k, max_len) for _ in range(seq)]
99 sequences = list(set(sequences)) # Удаление дубликатов
100
101 # Ранжирование
102 embd = self.__det_mean(embds)
103 embeddings = self.sbert.encode(sequences)
104 similarities = [self.cosine_similarity(embd, emb) for emb in embeddings]
105 best_index = np.argmax(similarities)
106
107 return sequences[best_index]pip install sentence-transformers -q1from sentence_transformers import SentenceTransformer
2
3sbert = SentenceTransformer('FractalGPT/SbertDistil')
4generator = TextEmbdGenerator('FractalGPT/EmbedderDecoder', sbert)1embd = sbert.encode('там живут англичане')
2generator.generate_with_ranker([embd])>>> я бы его в Англию привез.1embd = sbert.encode('там живут немцы')
2generator.generate_with_ranker([embd], prompt = 'он всегда был в')>>> он всегда был в Германии1embd = sbert.encode('он сделает вывод на основе анализа ситуации')
2generator.generate_with_ranker(embd)>>> в процессе анализа ситуации необходимо выяснить:1embd = sbert.encode('машина') - sbert.encode('колеса') + sbert.encode('крылья')
2generator.generate_with_ranker([embd], 'это')>>> этот самолёт1embd = sbert.encode('полицейский - главный герой') + sbert.encode('Произошло ужасное событие в фильме')
2embd /= 2
3generator.generate_with_ranker([embd], 'Собеседование на')>>> Собеседование на роль главного героя фильма — молодого лейтенанта полиции — происходит в доме1embd_1 = sbert.encode('полицейский - главный герой')
2embd_2 = sbert.encode('Произошло событие в фильме')
3
4generator.generate_with_ranker([embd_1, embd_2], 'В ')>>> В этом фильме главный герой - полицейский.1embd_1 = sbert.encode('полицейский - главный герой')
2embd_2 = sbert.encode('Произошло событие в фильме')
3
4generator.generate_with_ranker([embd_1], 'Это', [embd_2])>>> Это полицейский, который в полицейском участке снимается в фильме.1embd = sbert.encode('радиоприемник')
2ans = vector_answer(embd, 'Как это устроено?')
3print(ans)>>> Радиоволны распространяются в воздухе, создавая электромагнитное поле, которое может быть использовано для передачи информации.