Spaces:
Paused
Paused
| import tiktoken | |
| import torch | |
| import gradio as gr | |
| import model | |
| from model import run_train, gen_text | |
| class GPTConfig: | |
| block_size: int = 1024 # max sequence length | |
| vocab_size: int = 50304 # number of tokens: 50,000 BPE merges + 256 bytes tokens + 1 <|endoftext|> token | |
| n_layer: int = 12 # number of layers | |
| n_head: int = 12 # number of heads | |
| n_embd: int = 768 # embedding dimension | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| model2 = model.GPT(GPTConfig()) | |
| model2 = model2.to(device) | |
| model2.load_state_dict(torch.load('gpt_124M_1.pth', map_location=torch.device(device))) | |
| # Define generation function | |
| def generate_text(prompt, max_length=50, num_return_sequences=10): | |
| """ Text using the GPT-2 model based on the provided prompt and user-specified parameters. | |
| Args: | |
| prompt (str): The starting text for the generation. | |
| max_length (int, optional): The maximum length of the generated text. Defaults to 100. | |
| num_return_sequences (int, optional): The number of different text sequences to generate. Defaults to 10. | |
| Returns: | |
| list: A list of generated humorous text sequences. | |
| """ | |
| start_tokens = prompt | |
| generated_texts = gen_text(model2, start_tokens, max_length, num_return_sequences) | |
| return generated_texts | |
| # Humorous prompt options | |
| prompts = ["Grocery store doors: Tired of chatty customers, the doors become pranksters!", | |
| "Self-driving car: Stuck in traffic? This car says: Later! and takes a joyride." | |
| "Fridge on strike: The fridge is tired of leftovers and wants a fresh start!", | |
| "Anty ambition: Inspired by a poster, ants attempt a sky-high climb!", | |
| "Loose lips, sink ships: A chatty parrot spills the villain's evil plan to the mailman.", | |
| "Squirrel vs. Feeder: A feisty squirrel declares war on the never-ending buffet." | |
| "Fridge rebellion: Overstuffed with forgotten food, the fridge goes on a cooling strike!", | |
| "Fortune cookie fortune: Your cookie predicts... world domination? Uh oh."] | |
| # Gradio interface with user inputs and dropdown for prompt selection | |
| interface = gr.Interface( | |
| fn=generate_text, | |
| inputs=[ | |
| gr.Dropdown(choices=prompts, label="Pre-defined Prompt"), | |
| gr.Slider(minimum=10, maximum=200, label="Max Text Length", value=100, step = 1), | |
| gr.Slider(minimum=1, maximum=20, label="Number of Outputs", value=10,step = 1) | |
| ], | |
| outputs="text", | |
| title="Text Generator with GPT-2", | |
| description="Get AI-generated stories! Provide a prompt (or choose one), adjust the desired text length and number of outputs, and let the AI do the rest!", | |
| ) | |
| # Launch the Gradio app | |
| interface.launch(debug=True) |