Shakespeare-YQG / app.py
NightPrince's picture
Create app.py
d3a1a00 verified
import gradio as gr
import tensorflow as tf
import numpy as np
from tensorflow.keras.preprocessing.sequence import pad_sequences
import json
import os
# Initialize global variables
model = None
tokenizer = None
max_len_seq = None
def load_model_artifacts():
global model, tokenizer, max_len_seq
# Load model directly from root directory
model = tf.keras.models.load_model('shakespeare_model.h5')
# Load tokenizer from root
with open('tokenizer.json', 'r') as f:
tokenizer = tf.keras.preprocessing.text.tokenizer_from_json(f.read())
# Load config from root
with open('config.json', 'r') as f:
config = json.load(f)
max_len_seq = config['max_len_seq']
def generate_shakespeare_quote(seed_text, num_words):
"""Generate Shakespeare-style text from a seed text"""
if model is None:
load_model_artifacts()
try:
for _ in range(int(num_words)):
# Convert the seed text to sequences
token_list = tokenizer.texts_to_sequences([seed_text])[0]
# Pad the sequences
token_list = pad_sequences([token_list], maxlen=max_len_seq-1, padding='pre')
# Predict the next word
predicted = model.predict(token_list, verbose=0)
# Get the word with highest probability
predicted_word = tokenizer.index_word[np.argmax(predicted, axis=-1).item()]
# Add the predicted word to the seed text
seed_text += " " + predicted_word
return seed_text
except Exception as e:
return f"Error generating text: {str(e)}"
# Create the Gradio interface
iface = gr.Interface(
fn=generate_shakespeare_quote,
inputs=[
gr.Textbox(
label="Enter your seed text",
placeholder="Start your quote here...",
value="to be or"
),
gr.Slider(
minimum=1,
maximum=50,
value=10,
step=1,
label="Number of words to generate"
)
],
outputs=gr.Textbox(label="Generated Quote"),
title="Shakespeare Quote Generator",
description="""Generate Shakespeare-style quotes using AI!
Enter a seed text and choose how many words you want to generate.
The model will continue your text in Shakespeare's style.""",
examples=[
["to be or", 10],
["love is", 15],
["life is", 12],
["death be not", 10],
["shall i compare thee", 8]
],
theme=gr.themes.Base()
)
# Launch the app
if __name__ == "__main__":
iface.launch()