File size: 2,580 Bytes
a3b29ff |
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 65 66 |
import torch
import Train
import Tokenizer
def load_model(model_path, d_model, ffn_hidden, num_heads, drop_prob, num_layers):
model = Train.Transformer(d_model, ffn_hidden, num_heads, drop_prob, num_layers)
model.load_state_dict(torch.load(model_path))
model.eval()
return model
def prepare_input(input_text):
input_tokens = Tokenizer.tokenize_sequence(input_text)
# input_tokens = Tokenizer.pad_to_length(input_tokens, Train.max_sequence_length)
input_ids = torch.tensor(input_tokens).unsqueeze(0) # Add batch dimension
return input_ids
# def generate_output(model, input_ids):
# with torch.no_grad():
# output_logits = model(input_ids)
# predicted_token_ids = torch.argmax(output_logits, dim=-1)
# output_text = Tokenizer.detokenize_sequence(predicted_token_ids[0].tolist())
# return output_text
# def generate_output(model, input_ids, max_length, eos_token=Tokenizer.vocabulary.get('<EOS>')):
# with torch.no_grad(): # No need to track gradients during inference
# input_tensor = torch.tensor(input_ids)
# output_seq = []
#
# for i in range(50):
# output = model.generate(input_tensor)
# print(f'output.size(): {output.size()}')
# next_token = torch.argmax(output[0, i, :], dim=-1).item() # Take last token from sequence
# output_seq.append(next_token)
# if next_token == eos_token:
# break # Stop if EOS token is generated
# next_token_tensor = torch.tensor([[next_token]]).to(Train.device) # Convert and move to device
# input_tensor = torch.cat([input_tensor, next_token_tensor], dim=1) # Concatenate
# print(f'Generated tokens: {output_seq}')
# AI_response = Tokenizer.detokenize_sequence(output_seq)
# return AI_response
def generate_output(model, input_ids, max_length, eos_token=Tokenizer.vocabulary.get('<EOS>')):
out = model.generate(input_ids)
preds = torch.argmax(out, dim=-1)
output_tokens = []
for token in preds[0]:
output_tokens.append(token.item())
AI_response = Tokenizer.detokenize_sequence(output_tokens)
return AI_response
# Example usage
model_path = 'models/my_model.pt'
input_text = ''
model = load_model(model_path, Train.d_model, Train.ffn_hidden, Train.num_heads, Train.drop_prob, Train.num_layers)
model.to(Train.device)
input_ids = prepare_input(input_text)
output_text = generate_output(model, input_ids.to(Train.device), Train.max_sequence_length)
print("Generated Output:", output_text) |