Spaces:
Runtime error
Runtime error
File size: 2,042 Bytes
1f3cb28 |
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 |
import torch
from transformers import GPT2LMHeadModel, GPT2Tokenizer
def select_model():
while True:
print("\nAvailable GPT-2 Models:")
print("1. gpt2 (Small)")
print("2. gpt2-medium (Medium)")
print("3. gpt2-large (Large)")
print("4. gpt2-xl (Extra Large)")
choice = input("Select a model (1/2/3/4): ")
if choice == "1":
return "gpt2"
elif choice == "2":
return "gpt2-medium"
elif choice == "3":
return "gpt2-large"
elif choice == "4":
return "gpt2-xl"
else:
print("Invalid choice. Please select a valid model.")
def enhance_prompt(prompt, model_name, max_length=50, num_return_sequences=1):
# Load the selected pre-trained GPT-2 model and tokenizer
model = GPT2LMHeadModel.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
# Tokenize the prompt
input_ids = tokenizer.encode(prompt, return_tensors="pt")
# Generate text based on the prompt
output = model.generate(
input_ids,
max_length=max_length,
num_return_sequences=num_return_sequences,
no_repeat_ngram_size=2, # Avoid repetitive phrases
top_k=50, # Limit choices to top-k tokens
top_p=0.95, # Control diversity with nucleus sampling
temperature=0.7 # Adjusts the randomness of the output
)
# Decode and return the generated text
enhanced_prompts = [tokenizer.decode(output_item, skip_special_tokens=True) for output_item in output]
return enhanced_prompts
if __name__ == "__main__":
while True:
model_name = select_model()
prompt = input("Enter a prompt (or 'exit' to quit): ")
if prompt.lower() == "exit":
break
enhanced_prompts = enhance_prompt(prompt, model_name)
print("\nEnhanced Prompts:")
for idx, enhanced_prompt in enumerate(enhanced_prompts):
print(f"Enhanced Prompt {idx + 1}: {enhanced_prompt}\n")
|