Views
No views yet
1from transformers import GPT2LMHeadModel, GPT2Tokenizer
2
3# Load model and tokenizer
4model = GPT2LMHeadModel.from_pretrained("nesar2004/message-generator-gpt2")
5tokenizer = GPT2Tokenizer.from_pretrained("nesar2004/message-generator-gpt2")
6
7# Generate message
8category = "dating" # Can be: dating, flirty, love, cheesy
9input_text = f"<|{category}|><|message|>"
10input_ids = tokenizer.encode(input_text, return_tensors='pt')
11
12# Generate
13output = model.generate(
14 input_ids,
15 max_length=60,
16 num_return_sequences=1,
17 temperature=0.9,
18 top_p=0.95,
19 do_sample=True,
20 pad_token_id=tokenizer.pad_token_id,
21 eos_token_id=tokenizer.eos_token_id
22)
23
24# Decode and clean up the message
25message = tokenizer.decode(output[0], skip_special_tokens=True)
26message = message.replace("<|message|>", "").strip()
27message = message.replace(f"<|{category}|>", "").strip()
28print(message)