File size: 1,221 Bytes
a15e210 09e89d7 a15e210 28e3315 a15e210 28e3315 a15e210 |
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 |
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch
# Load the model and tokenizer
model_path = 'model_data/finetuned_gpt'
tokenizer = GPT2Tokenizer.from_pretrained(model_path)
model = GPT2LMHeadModel.from_pretrained(model_path)
# Move model to GPU if available
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
def generate_text(prompt_text, length, temperature, beams):
encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False, return_tensors="pt")
encoded_prompt = encoded_prompt.to(device)
output_sequences = model.generate(
input_ids=encoded_prompt,
max_length=length,
temperature=temperature,
top_k=20,
top_p=0.9,
repetition_penalty=1.2,
do_sample=True,
num_return_sequences=beams,
)
# Decode the generated text
generated_sequence = output_sequences[0].tolist()
text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True)
# Remove the prompt from the generated text
text = text[len(tokenizer.decode(encoded_prompt[0], clean_up_tokenization_spaces=True)) :]
return text.strip()
# Streamlit interface
|