TheGiggleGrid / app.py
SeemG's picture
Update app.py
54c080c verified
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)