-
Notifications
You must be signed in to change notification settings - Fork 0
/
runner.py
35 lines (29 loc) · 1.3 KB
/
runner.py
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
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer
async def generate(prefix, max_length=800, top_k=5, model_dir='fine_tuned_model'):
"""Generate the idea based on inputed prefix"""
# Load the fine-tuned model and tokenizer
model = GPT2LMHeadModel.from_pretrained(model_dir)
tokenizer = GPT2Tokenizer.from_pretrained(model_dir)
# Encode the prefix
input_ids = tokenizer.encode(prefix, add_special_tokens=True, return_tensors='pt')
# Generate text based on the prefix
model.eval()
model.config.pad_token_id = tokenizer.eos_token_id
model.config.eos_token_id = tokenizer.eos_token_id
model.config.vocab_size = model.config.vocab_size + len(tokenizer.get_added_vocab())
model.resize_token_embeddings(len(tokenizer))
# model, input_ids = model.to(torch.device("cuda")), input_ids.to(torch.device("cuda"))
with torch.no_grad():
output = model.generate(
input_ids,
max_length=max_length,
num_return_sequences=1,
top_k=top_k,
do_sample=True,
pad_token_id=tokenizer.eos_token_id,
eos_token_id=tokenizer.eos_token_id
)
# Decode and return the generated text
generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
return generated_text