File size: 2,557 Bytes
cb69e53 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
import torch
from transformers import T5Tokenizer
from model import GPT
class Inference:
def __init__(self, model_path, tokenizer_path, device='cuda' if torch.cuda.is_available() else 'cpu'):
self.device = device
self.tokenizer = T5Tokenizer.from_pretrained(tokenizer_path)
self.model = GPT(
vocab_size=self.tokenizer.vocab_size,
embed_size=1500,
num_layers=20,
heads=20,
expansion_factor=4,
dropout=0.1,
max_length=1024
)
self.model.load_state_dict(torch.load(model_path, map_location=self.device))
self.model.to(self.device)
self.model.eval()
def predict(self, text, max_length=100):
input_ids = self.tokenizer.encode(text, return_tensors='pt').to(self.device)
generated_tokens = set(input_ids[0].tolist())
with torch.no_grad():
for _ in range(max_length):
outputs = self.model(input_ids)
logits = outputs[:, -1, :] / 1.0 # temperature = 1.0
for token_id in generated_tokens:
logits[0, token_id] /= 1.5 # repetition_penalty = 1.5
filtered_logits = top_k_top_p_filtering(logits, top_k=50, top_p=0.9)
probs = torch.softmax(filtered_logits, dim=-1)
next_token_id = torch.multinomial(probs, 1)
next_token_id = next_token_id.squeeze(-1).unsqueeze(0)
input_ids = torch.cat([input_ids, next_token_id], dim=1)
generated_tokens.add(next_token_id.item())
if next_token_id.item() == self.tokenizer.eos_token_id:
break
return self.tokenizer.decode(input_ids[0], skip_special_tokens=True)
def top_k_top_p_filtering(logits, top_k=0, top_p=0.9, filter_value=-float('Inf')):
top_k = min(top_k, logits.size(-1))
if top_k > 0:
indices_to_remove = logits < torch.topk(logits, top_k).values[:, -1, None]
logits[indices_to_remove] = filter_value
sorted_logits, sorted_indices = torch.sort(logits, descending=True)
cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
sorted_indices_to_remove = cumulative_probs > top_p
sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
sorted_indices_to_remove[..., 0] = 0
indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
logits[indices_to_remove] = filter_value
return logits
|