Files changed (1) hide show
  1. README.md +11 -0
README.md ADDED
@@ -0,0 +1,11 @@
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import GPT2LMHeadModel, GPT2Tokenizer
2
+
3
+ # Load the pre-trained WormGPT model and tokenizer
4
+ model = GPT2LMHeadModel.from_pretrained("wormgpt")
5
+ tokenizer = GPT2Tokenizer.from_pretrained("wormgpt")
6
+
7
+ def generate_text(prompt, max_length=50):
8
+ input_ids = tokenizer.encode(prompt, return_tensors="pt")
9
+ output = model.generate(input_ids, max_length=max_length, num_return_sequences=1)
10
+ generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
11
+ return generated_text